Repository: huggingface/datatrove (language: python)

Issue:
:sparkles: ExclusionWriter for filters
- added a writer that can save samples removed by filters
- made it so iterables can be passed as a pipeline step
- small bugfix associated with pickling


Discussion / hints:

Fix patch:
diff --git a/examples/other sources/tokenize_the_pile.py b/examples/other sources/tokenize_the_pile.py
index 72e364de..6bde05ec 100644
--- a/examples/other sources/tokenize_the_pile.py
+++ b/examples/other sources/tokenize_the_pile.py
@@ -1,36 +1,58 @@
-import os.path
-
 from datatrove.executor.base import PipelineExecutor
 from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.io import LocalOutputDataFolder, S3InputDataFolder
+from datatrove.io import LocalInputDataFolder, LocalOutputDataFolder, S3InputDataFolder, S3OutputDataFolder
 from datatrove.pipeline.readers import JsonlReader
+from datatrove.pipeline.tokens.merger import DocumentTokenizerMerger
 from datatrove.pipeline.tokens.tokenizer import DocumentTokenizer
 
 
-def format_adapter(d: dict, path: str, li: int):
-    return {
-        "content": d["text"],
-        "data_id": f"{os.path.splitext(os.path.basename(path))[0]}_{li}",
-        "metadata": d["meta"],
-    }
-
-
 pipeline = [
     JsonlReader(
         S3InputDataFolder("s3://bigcode-experiments/the-pile-sharded/", stream=True),
-        gzip=False,
-        adapter=format_adapter,
+        content_key="text",
+        limit=100,
     ),
     DocumentTokenizer(
-        LocalOutputDataFolder(path="/fsx/guilherme/the-pile/tokenized"), save_filename="the_pile_tokenized"
+        LocalOutputDataFolder(path="/fsx/guilherme/tests/piletokenized-test/tokenized"),
+        save_filename="the_pile_tokenized",
     ),
 ]
 
 executor: PipelineExecutor = SlurmPipelineExecutor(
     pipeline=pipeline,
-    tasks=20,
-    time="12:00:00",
-    partition="dev-cluster",
-    logging_dir="/fsx/guilherme/logs/tokenize_the_pile",
+    tasks=2,
+    workers=1,
+    skip_completed=False,
+    logging_dir=S3OutputDataFolder(
+        "s3://extreme-scale-dp-temp/logs/tests/piletokenized/tokenized",
+        local_path="/fsx/guilherme/tests/piletokenized-test/logs",
+        cleanup=False,
+    ),
+    slurm_logs_folder="/fsx/guilherme/tests/piletokenized-test/slurm_logs",
+    partition="production-cluster",
+    time="01:00:00",
 )
 executor.run()
+
+executor2: PipelineExecutor = SlurmPipelineExecutor(
+    pipeline=[
+        DocumentTokenizerMerger(
+            input_folder=LocalInputDataFolder(path="/fsx/guilherme/tests/piletokenized-test/tokenized"),
+            output_folder=LocalOutputDataFolder(path="/fsx/guilherme/tests/piletokenized-test/merged"),
+            save_filename="hello",
+            max_tokens_per_file=100000,
+        )
+    ],
+    depends=executor,
+    tasks=1,
+    skip_completed=False,
+    logging_dir=S3OutputDataFolder(
+        "s3://extreme-scale-dp-temp/logs/tests/piletokenized/merged",
+        local_path="/fsx/guilherme/tests/piletokenized-test/logs2",
+        cleanup=False,
+    ),
+    slurm_logs_folder="/fsx/guilherme/tests/piletokenized-test/slurm_logs2",
+    partition="production-cluster",
+    time="01:00:00",
+)
+executor2.run()
diff --git a/setup.py b/setup.py
index 8f0c88c1..12cf5e8e 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,7 @@
     "pandas==2.0.3",
     "backoff==2.2.1",
     "fsspec==2023.9.2",
+    "humanize==4.8.0",
 ]
 
 EXTRAS = {
@@ -68,6 +69,7 @@
         "console_scripts": [
             "check_dataset=datatrove.tools.check_dataset:main",
             "merge_stats=datatrove.tools.merge_stats:main",
+            "launch_pickled_pipeline=datatrove.tools.launch_pickled_pipeline:main",
         ]
     },
 )
diff --git a/src/datatrove/executor/base.py b/src/datatrove/executor/base.py
index c46e9c36..e197391c 100644
--- a/src/datatrove/executor/base.py
+++ b/src/datatrove/executor/base.py
@@ -1,3 +1,5 @@
+import dataclasses
+import json
 from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Sequence
@@ -5,18 +7,25 @@
 
 from loguru import logger
 
+from datatrove.io import BaseOutputDataFolder, LocalOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
+from datatrove.utils.logging import add_task_logger, close_task_logger, get_random_str, get_timestamp, log_pipeline
 from datatrove.utils.stats import PipelineStats
 
 
 class PipelineExecutor(ABC):
     @abstractmethod
-    def __init__(self, pipeline: list[PipelineStep | Callable], save_stats: bool = False):
+    def __init__(
+        self,
+        pipeline: list[PipelineStep | Callable],
+        logging_dir: BaseOutputDataFolder = None,
+        skip_completed: bool = True,
+    ):
         self.pipeline: list[PipelineStep | Callable] = pipeline
-        self.save_stats = save_stats
-
-        pipeline = "\n".join([pipe.__repr__() if callable(pipe) else "Sequence..." for pipe in self.pipeline])
-        print(f"--- 🛠️PIPELINE 🛠\n{pipeline}")
+        self.logging_dir = (
+            logging_dir if logging_dir else LocalOutputDataFolder(f"logs/{get_timestamp()}_{get_random_str()}")
+        )
+        self.skip_completed = skip_completed
 
     @abstractmethod
     def run(self):
@@ -27,20 +36,62 @@ def run(self):
     def world_size(self):
         return 0
 
-    def _run_for_rank(self, rank: int) -> PipelineStats:
-        logger.info(f"Launching pipeline for {rank=}")
-        # pipe data from one step to the next
-        pipelined_data = None
-        for pipeline_step in self.pipeline:
-            if callable(pipeline_step):
-                pipelined_data = pipeline_step(pipelined_data, rank, self.world_size)
-            elif isinstance(pipeline_step, Sequence) and not isinstance(pipeline_step, str):
-                pipelined_data = pipeline_step
-            else:
-                raise ValueError
-        if pipelined_data:
-            deque(pipelined_data, maxlen=0)
-        logger.info(f"Processing done for {rank=}")
-        return PipelineStats(
-            [pipeline_step.stats for pipeline_step in self.pipeline if isinstance(pipeline_step, PipelineStep)]
-        )
+    def _run_for_rank(self, rank: int, local_rank: int = 0) -> PipelineStats:
+        if self.is_rank_completed(rank):
+            logger.info(f"Skipping {rank=} as it has already been completed.")
+            return PipelineStats()
+        add_task_logger(self.logging_dir, rank, local_rank)
+        log_pipeline(self.pipeline)
+        try:
+            # pipe data from one step to the next
+            pipelined_data = None
+            for pipeline_step in self.pipeline:
+                if callable(pipeline_step):
+                    pipelined_data = pipeline_step(pipelined_data, rank, self.world_size)
+                elif isinstance(pipeline_step, Sequence) and not isinstance(pipeline_step, str):
+                    pipelined_data = pipeline_step
+                else:
+                    raise ValueError
+            if pipelined_data:
+                deque(pipelined_data, maxlen=0)
+
+            logger.success(f"Processing done for {rank=}")
+
+            # stats
+            stats = PipelineStats(self.pipeline)
+            stats.save_to_disk(self.logging_dir.open(f"stats/{rank:05d}.json"))
+            logger.info(stats.get_repr(f"Task {rank}"))
+            # completed
+            self.mark_rank_as_completed(rank)
+        except Exception as e:
+            logger.exception(e)
+            raise e
+        finally:
+            close_task_logger(self.logging_dir, rank)
+        return stats
+
+    def is_rank_completed(self, rank: int):
+        return self.skip_completed and self.logging_dir.to_input_folder().file_exists(f"completions/{rank:05d}")
+
+    def mark_rank_as_completed(self, rank: int):
+        self.logging_dir.open(f"completions/{rank:05d}").close()
+
+    def to_json(self, indent=4):
+        data = self.__dict__
+        data["pipeline"] = [{a: b for a, b in x.__dict__.items() if a != "stats"} for x in data["pipeline"]]
+        return json.dumps(data, indent=indent)
+
+    def save_executor_as_json(self, indent: int = 4):
+        with self.logging_dir.open("executor.json") as f:
+            json.dump(self, f, cls=ExecutorJSONEncoder, indent=indent)
+
+
+class ExecutorJSONEncoder(json.JSONEncoder):
+    def default(self, o):
+        if dataclasses.is_dataclass(o):
+            return dataclasses.asdict(o)
+        if isinstance(o, PipelineExecutor):
+            return o.__dict__
+        if isinstance(o, PipelineStep):
+            return {a: b for a, b in o.__dict__.items() if a != "stats"}
+        return str(o)
diff --git a/src/datatrove/executor/local.py b/src/datatrove/executor/local.py
index 6ccc8b5b..0c18032b 100644
--- a/src/datatrove/executor/local.py
+++ b/src/datatrove/executor/local.py
@@ -1,56 +1,117 @@
+import itertools
 from copy import deepcopy
+from multiprocessing import Value
+from typing import Callable
 
 import multiprocess.pool
-from multiprocess import Semaphore
+from loguru import logger
+from multiprocess import Queue, Semaphore
 
 from datatrove.executor.base import PipelineExecutor
+from datatrove.io import BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
+from datatrove.utils.stats import PipelineStats
 
 
-download_semaphore, upload_semaphore = None, None
+# multiprocessing vars
+download_semaphore, upload_semaphore, ranks_queue, completed = None, None, None, None
 
 
-def init_pool_processes(dl_sem, up_sem):
-    global download_semaphore, upload_semaphore
+def init_pool_processes(dl_sem, up_sem, ranks_q, completed_counter):
+    global download_semaphore, upload_semaphore, ranks_queue, completed
     download_semaphore = dl_sem
     upload_semaphore = up_sem
+    ranks_queue = ranks_q
+    completed = completed_counter
 
 
 class LocalPipelineExecutor(PipelineExecutor):
     def __init__(
         self,
+        pipeline: list[PipelineStep | Callable],
         tasks: int = 1,
         workers: int = -1,
         max_concurrent_uploads: int = 20,
         max_concurrent_downloads: int = 50,
-        **kwargs,
+        logging_dir: BaseOutputDataFolder = None,
+        skip_completed: bool = True,
     ):
-        super().__init__(**kwargs)
+        """Execute a pipeline on a slurm cluster
+
+        Args:
+            pipeline: a list of PipelineStep and/or custom functions
+                with arguments (data: DocumentsPipeline, rank: int,
+                world_size: int)
+            tasks: total number of tasks to run the pipeline on
+            workers: how many tasks to run simultaneously. -1 for no
+                limit
+            max_concurrent_uploads: limit the number of files that may
+                be uploaded simultaneously to avoid rate limits
+            max_concurrent_downloads: limit the number of files that may
+                be downloaded simultaneously to avoid rate limits
+            logging_dir: where to save logs, stats, etc. Should be an
+                OutputDataFolder
+            skip_completed: whether to skip tasks that were completed in
+                previous runs. default: True
+        """
+        super().__init__(pipeline, logging_dir, skip_completed)
         self.tasks = tasks
         self.workers = workers if workers != -1 else tasks
         self.max_concurrent_uploads = max_concurrent_uploads
         self.max_concurrent_downloads = max_concurrent_downloads
 
-    def _run_for_rank(self, rank: int):
+    def _run_for_rank(self, rank: int, local_rank: int = -1):
         if self.workers > 1:
             for pipeline_step in self.pipeline:
                 if isinstance(pipeline_step, PipelineStep):
                     pipeline_step.set_up_dl_locks(download_semaphore, upload_semaphore)
-        return super()._run_for_rank(rank)
+        local_rank = ranks_queue.get()
+        try:
+            return super()._run_for_rank(rank, local_rank)
+        finally:
+            if completed:
+                with completed.get_lock():
+                    completed.value += 1
+                    logger.info(f"{completed.value}/{self.world_size} tasks completed.")
+            ranks_queue.put(local_rank)  # free up used rank
 
     def run(self):
+        if all(map(self.is_rank_completed, range(self.tasks))):
+            logger.info(f"Not doing anything as all {self.tasks} tasks have already been completed.")
+            return
+
+        self.save_executor_as_json()
+        ranks_q = Queue()
+        for i in range(self.workers):
+            ranks_q.put(i)
+
+        ranks_to_run = list(itertools.filterfalse(self.is_rank_completed, range(self.tasks)))
+        if (skipped := self.tasks - len(ranks_to_run)) > 0:
+            logger.info(f"Skipping {skipped} already completed tasks")
+
         if self.workers == 1:
+            global ranks_queue
+            ranks_queue = ranks_q
             pipeline = self.pipeline
             stats = []
-            for rank in range(self.tasks):
+            for rank in ranks_to_run:
                 self.pipeline = deepcopy(pipeline)
                 stats.append(self._run_for_rank(rank))
         else:
             dl_sem = Semaphore(self.max_concurrent_downloads)
             up_sem = Semaphore(self.max_concurrent_uploads)
-            with multiprocess.Pool(self.workers, initializer=init_pool_processes, initargs=(dl_sem, up_sem)) as pool:
-                stats = list(pool.map(self._run_for_rank, range(self.tasks)))
-        return sum(stats)
+            completed_counter = Value("i", skipped)
+
+            with multiprocess.Pool(
+                self.workers, initializer=init_pool_processes, initargs=(dl_sem, up_sem, ranks_q, completed_counter)
+            ) as pool:
+                stats = list(pool.imap_unordered(self._run_for_rank, ranks_to_run))
+        # merged stats
+        stats = sum(stats, start=PipelineStats())
+        stats.save_to_disk(self.logging_dir.open("stats.json"))
+        logger.success(stats.get_repr(f"All {self.tasks} tasks"))
+        self.logging_dir.close()
+        return stats
 
     @property
     def world_size(self):
diff --git a/src/datatrove/executor/slurm.py b/src/datatrove/executor/slurm.py
index 60c88366..2f2d9014 100644
--- a/src/datatrove/executor/slurm.py
+++ b/src/datatrove/executor/slurm.py
@@ -1,27 +1,33 @@
 from __future__ import annotations
 
+import dataclasses
 import os
 import subprocess
 import tempfile
 import textwrap
+from copy import deepcopy
+from typing import Callable
 
 import dill
+from dill import CONTENTS_FMODE
 from loguru import logger
 
 from datatrove.executor.base import PipelineExecutor
+from datatrove.io import BaseOutputDataFolder, S3OutputDataFolder
+from datatrove.pipeline.base import PipelineStep
+from datatrove.utils.logging import get_random_str, get_timestamp
 
 
 class SlurmPipelineExecutor(PipelineExecutor):
-    """
-    Executor to run pipelines on Slurm.
+    """Executor to run pipelines on Slurm.
     Creates and calls a sbatch launch script.
     """
 
     def __init__(
         self,
+        pipeline: list[PipelineStep | Callable],
         tasks: int,
         time: str,
-        logging_dir: str,
         partition: str,
         cpus_per_task: int = 1,
         mem_per_cpu_gb: int = 2,
@@ -33,26 +39,54 @@ def __init__(
         sbatch_args: dict | None = None,
         max_array_size: int = 1001,
         depends: SlurmPipelineExecutor | None = None,
-        **kwargs,
+        logging_dir: BaseOutputDataFolder = None,
+        skip_completed: bool = True,
+        slurm_logs_folder: str = None,
     ):
+        """Execute a pipeline on a slurm cluster
+
+        Args:
+            pipeline: a list of PipelineStep and/or custom functions
+                with arguments (data: DocumentsPipeline, rank: int,
+                world_size: int)
+            tasks: total number of tasks to run the pipeline on
+            time: slurm time limit
+            partition: slurm partition
+            cpus_per_task: how many cpus to give each task. should be 1
+                except when you need to give each task more memory
+            mem_per_cpu_gb: slurm option. use in conjunction with the
+                above option to increase max memory
+            workers: how many tasks to run simultaneously. -1 for no
+                limit
+            job_name: slurm job name
+            env_command: command to activate a python environment, if
+                needed
+            condaenv: name of a conda environment to activate
+            venv_path: path to a python venv to activate
+            sbatch_args: dictionary with additional arguments to pass to
+                sbatch
+            max_array_size: the limit of tasks in a task array job on
+                your slurm cluster or -1 if none. if
+                tasks>max_array_size, multiple task array jobs will be
+                launched
+            depends: another SlurmPipelineExecutor that should run
+                before this one
+            logging_dir: where to save logs, stats, etc. Should be an
+                OutputDataFolder
+            skip_completed: whether to skip tasks that were completed in
+                previous runs. default: True
+            slurm_logs_folder: where to store the raw slurm log files.
+                must be a local path default:
+                slurm_logs/$job_name/$timestamp_$randomstring
         """
-        :param tasks: total number of tasks to run
-        :param time: time limit, passed to slurm
-        :param logging_dir: directory where log files and the launch script should be saved
-        :param cpus_per_task: how many cpus per task
-        :param job_name: slurm job name
-        :param condaenv: name of a conda environment to activate before starting the job
-        :param venv_path: path to a virtual environment to activate, if not using conda
-        :param sbatch_args: a dictionary of other SBATCH arguments for the launch script
-        :param kwargs:
-        """
-        super().__init__(**kwargs)
+        if isinstance(logging_dir, S3OutputDataFolder):
+            logging_dir.cleanup = False  # if the files are removed from disk job launch will fail
+        super().__init__(pipeline, logging_dir, skip_completed)
         self.tasks = tasks
         self.workers = workers
         self.partition = partition
         self.cpus_per_task = cpus_per_task
         self.mem_per_cpu_gb = mem_per_cpu_gb
-        self.logging_dir = logging_dir
         self.time = time
         self.job_name = job_name
         self.env_command = env_command
@@ -64,37 +98,38 @@ def __init__(
         self.job_ids = []
         self.depends_job_ids = []
         self.launched = False
+        self.slurm_logs_folder = (
+            slurm_logs_folder
+            if slurm_logs_folder
+            else f"slurm_logs/{self.job_name}/{get_timestamp()}_{get_random_str()}"
+        )
 
     def run(self):
         if "SLURM_ARRAY_TASK_ID" in os.environ:
             rank = int(os.environ["SLURM_ARRAY_TASK_ID"]) + self.max_array_size * int(os.environ.get("RUN_OFFSET", 0))
             if rank >= self.world_size:
                 return
-            if self.is_rank_completed(rank):
-                logger.info(f"Skipping {rank=} as it has already been completed.")
-                return
-            stats = self._run_for_rank(rank)
-            stats.save_to_disk(os.path.join(self.logging_dir, "stats", f"{rank:05d}.json"))
-            self.mark_rank_as_completed(rank)
+            self._run_for_rank(rank)
+            self.logging_dir.close()  # make sure everything is properly saved (logs etc)
         else:
             self.launch_job()
 
     def launch_merge_stats(self):
-        with tempfile.NamedTemporaryFile("w") as f:
-            f.write(
-                self.get_launch_file(
-                    {
-                        **self.sbatch_args,
-                        "cpus-per-task": 1,
-                        "mem-per-cpu": "1G",
-                        "array": "0",
-                        "dependency": f"afterok:{','.join(self.job_ids)}",
-                    },
-                    f'merge_stats {os.path.join(self.logging_dir, "stats")} --output {os.path.join(self.logging_dir, "stats.json")}',
-                )
+        stats_json_file = self.logging_dir.create_new_file("stats.json")
+        # dump outputfile
+        options = [f"{k}={v}" for k, v in dataclasses.asdict(stats_json_file).items() if not k.startswith("_")]
+        launch_slurm_job(
+            self.get_launch_file_contents(
+                {
+                    **self.sbatch_args,
+                    "cpus-per-task": 1,
+                    "mem-per-cpu": "1G",
+                    "array": "0",
+                    "dependency": f"afterok:{','.join(self.job_ids)}",
+                },
+                f'merge_stats {os.path.join(self.logging_dir.path, "stats")} {" ".join(options)}',
             )
-            f.flush()
-            subprocess.check_output(["sbatch", f.name])
+        )
 
     @property
     def dependency(self):
@@ -106,11 +141,6 @@ def dependency(self):
         return ",".join(dependency)
 
     def launch_job(self):
-        launch_script_path = os.path.join(self.logging_dir, "launch_script.slurm")
-        os.makedirs(os.path.join(self.logging_dir, "stats"), exist_ok=True)
-        os.makedirs(os.path.join(self.logging_dir, "logs"), exist_ok=True)
-        os.makedirs(os.path.join(self.logging_dir, "completions"), exist_ok=True)
-
         assert not self.depends or (
             isinstance(self.depends, SlurmPipelineExecutor)
         ), "depends= must be a SlurmPipelineExecutor"
@@ -126,25 +156,32 @@ def launch_job(self):
             self.launched = True
             return
 
-        logger.info(f'Launching Slurm job {self.job_name} with launch script "{launch_script_path}"')
+        executor = deepcopy(self)
 
         # pickle
-        with open(os.path.join(self.logging_dir, "executor.pik"), "wb") as f:
-            dill.dump(self, f)
+        with self.logging_dir.open("executor.pik", "wb") as executor_f:
+            dill.dump(executor, executor_f, fmode=CONTENTS_FMODE)
+        self.save_executor_as_json()
 
-        with open(launch_script_path, "w") as f:
-            f.write(self.launch_file)
+        launch_file_contents = self.get_launch_file_contents(
+            self.sbatch_args,
+            f"srun -l launch_pickled_pipeline {executor_f.path}",
+        )
+        with self.logging_dir.open("launch_script.slurm") as launchscript_f:
+            launchscript_f.write(launch_file_contents)
+        logger.info(f'Launching Slurm job {self.job_name} with launch script "{launchscript_f.path}"')
 
         self.job_ids = []
         while len(self.job_ids) * self.max_array < self.tasks:
-            args = ["sbatch", f"--export=ALL,RUN_OFFSET={len(self.job_ids)}"]
+            args = [f"--export=ALL,RUN_OFFSET={len(self.job_ids)}"]
             if self.dependency:
                 args.append(f"--dependency={self.dependency}")
-            output = subprocess.check_output(args + [launch_script_path]).decode("utf-8")
-            self.job_ids.append(output.split()[-1])
+            launched_id = launch_slurm_job(launch_file_contents, *args)
+            self.job_ids.append(launched_id)
         self.launched = True
         logger.info(f"Slurm job launched successfully with id(s)={','.join(self.job_ids)}.")
         self.launch_merge_stats()
+        self.logging_dir.close()
 
     @property
     def max_array(self) -> int:
@@ -152,7 +189,8 @@ def max_array(self) -> int:
 
     @property
     def sbatch_args(self) -> dict:
-        slurm_logfile = os.path.join(self.logging_dir, "logs", "%j.out")
+        os.makedirs(self.slurm_logs_folder, exist_ok=True)
+        slurm_logfile = os.path.join(self.slurm_logs_folder, "%j.out")
         return {
             "cpus-per-task": self.cpus_per_task,
             "mem-per-cpu": f"{self.mem_per_cpu_gb}G",
@@ -165,20 +203,7 @@ def sbatch_args(self) -> dict:
             **self._sbatch_args,
         }
 
-    @property
-    def launch_file(self):
-        return self.get_launch_file(
-            self.sbatch_args,
-            f'srun -l python -u -c "import dill;dill.load(open(\'{os.path.join(self.logging_dir, "executor.pik")}\', \'rb\')).run()"',
-        )
-
-    def is_rank_completed(self, rank: int):
-        return os.path.exists(os.path.join(self.logging_dir, "completions", f"{rank:05d}"))
-
-    def mark_rank_as_completed(self, rank: int):
-        open(os.path.join(self.logging_dir, "completions", f"{rank:05d}"), "w").close()
-
-    def get_launch_file(self, sbatch_args: dict, run_script: str):
+    def get_launch_file_contents(self, sbatch_args: dict, run_script: str):
         args = "\n".join([f"#SBATCH --{k}={v}" for k, v in sbatch_args.items()])
 
         env_command = (
@@ -209,3 +234,10 @@ def get_launch_file(self, sbatch_args: dict, run_script: str):
     @property
     def world_size(self):
         return self.tasks
+
+
+def launch_slurm_job(launch_file_contents, *args):
+    with tempfile.NamedTemporaryFile("w") as f:
+        f.write(launch_file_contents)
+        f.flush()
+        return subprocess.check_output(["sbatch", *args, f.name]).decode("utf-8").split()[-1]
diff --git a/src/datatrove/io/__init__.py b/src/datatrove/io/__init__.py
index be970def..0ba94486 100644
--- a/src/datatrove/io/__init__.py
+++ b/src/datatrove/io/__init__.py
@@ -1,4 +1,4 @@
-from .base import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile, OutputDataFile
+from .base import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFile, BaseOutputDataFolder
 from .fsspec import FSSpecInputDataFile, FSSpecInputDataFolder, FSSpecOutputDataFile, FSSpecOutputDataFolder
-from .local import LocalInputDataFolder, LocalOutputDataFolder
+from .local import LocalInputDataFile, LocalInputDataFolder, LocalOutputDataFile, LocalOutputDataFolder
 from .s3 import S3InputDataFile, S3InputDataFolder, S3OutputDataFile, S3OutputDataFolder
diff --git a/src/datatrove/io/base.py b/src/datatrove/io/base.py
index 608e371a..b39c0ca0 100644
--- a/src/datatrove/io/base.py
+++ b/src/datatrove/io/base.py
@@ -1,5 +1,4 @@
 import contextlib
-import gzip as gzip_lib
 import os
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
@@ -9,24 +8,72 @@
 from io import TextIOWrapper
 from typing import Literal
 
+import numpy as np
 import zstandard
 from loguru import logger
+from multiprocess.synchronize import SemLock
 
 from datatrove.io.utils.fsspec import valid_fsspec_path
 
 
 @dataclass
-class InputDataFile:
+class BaseInputDataFile:
+    """Represents an individual input file that can be read from.
+    Supports compression formats and/or opening in binary mode.
+
+    Args:
+        path (str): complete path to this input file
+        relative_path (str): path relative to the input data folder this file belongs to
+    """
+
     path: str
     relative_path: str
+    _lock: SemLock | contextlib.nullcontext = field(default_factory=contextlib.nullcontext)
+
+    @classmethod
+    def from_path(cls, path: str, **kwargs):
+        """InputDataFile factory: get the correct instance from a path. Additional arguments will be passed to the
+        constructor of the matched implementation.
+
+        Args:
+            path: the full path to match
+            **kwargs: any additional argument passed to the matched
+                implementation
+
+        Returns:
+
+        """
+        from datatrove.io import FSSpecInputDataFile, LocalInputDataFile, S3InputDataFile
+
+        if "relative_path" not in kwargs:
+            kwargs["relative_path"] = os.path.basename(path)
+
+        if path.startswith("s3://"):
+            return S3InputDataFile(path, **kwargs)
+        elif valid_fsspec_path(path):
+            return FSSpecInputDataFile(path, **kwargs)
+        return LocalInputDataFile(path, **kwargs)
 
     @contextmanager
+    @abstractmethod
     def open_binary(self):
-        with open(self.path, mode="rb") as f:
-            yield f
+        """Main method to overwrite. Should return a stream open in binary mode
+
+        Returns:
+            binary stream
+        """
+        raise NotImplementedError
 
     @contextmanager
     def open_gzip(self, binary=False):
+        """Read a gzip compressed file directly
+
+        Args:
+            binary: optionally open in binary mode
+
+        Returns:
+            stream
+        """
         with self.open_binary() as fo:
             with GzipFile(mode="r" if not binary else "rb", fileobj=fo) as gf:
                 if binary:
@@ -37,6 +84,14 @@ def open_gzip(self, binary=False):
 
     @contextmanager
     def open_zst(self, binary=False):
+        """Read a zst compressed file directly
+
+        Args:
+            binary
+
+        Returns:
+            stream
+        """
         with self.open_binary() as fo:
             dctx = zstandard.ZstdDecompressor(max_window_size=2**31)
             stream_reader = dctx.stream_reader(fo)
@@ -47,7 +102,18 @@ def open_zst(self, binary=False):
                     yield f
 
     @contextmanager
-    def open(self, binary=False, compression: Literal["gzip", "zst"] | None = None):
+    def open(self, binary=False, compression: Literal["guess", "gzip", "zst"] | None = None):
+        """Main entrypoint for input files. Return a stream to read this file. Optionally accepts compression and/or binary mode.
+
+        Args:
+            binary
+            compression
+
+        Returns:
+
+        """
+        if compression == "guess":
+            compression = guess_compression(self.relative_path)
         match compression:
             case "gzip":
                 with self.open_gzip(binary) as f:
@@ -66,22 +132,35 @@ def open(self, binary=False, compression: Literal["gzip", "zst"] | None = None):
 
 @dataclass
 class BaseInputDataFolder(ABC):
-    """An input data folder
-
-    Args:
-        path (str): path to the folder
-        extension (str | list[str], optional): file extensions to filter. Defaults to None.
-        recursive (bool, optional): whether to search recursively. Defaults to True.
-        match_pattern (str, optional): pattern to match file names. Defaults to None.
+    """Base input data folder class. Specific implementations should override its relevant methods.
+
+        Args:
+    path (str): full path to the folder, respecting the format of specific implementations
+        (i.e. s3 paths should start with s3://)
+    extension (str | list[str], optional): file extensions to filter. Defaults to None.
+    recursive (bool, optional): whether to search recursively. Defaults to True.
+    match_pattern (str, optional): pattern to match file names. Defaults to None.
     """
 
     path: str
     extension: str | list[str] = None
     recursive: bool = True
     match_pattern: str = None
+    _lock: SemLock | contextlib.nullcontext = field(default_factory=contextlib.nullcontext)
 
     @classmethod
     def from_path(cls, path, **kwargs):
+        """InputDataFolder factory: get the correct instance from a path. Additional arguments will be passed to the
+        constructor of the matched implementation.
+
+        Args:
+            path: the full path to match
+            **kwargs: any additional argument passed to the matched
+                implementation
+
+        Returns:
+
+        """
         from datatrove.io import FSSpecInputDataFolder, LocalInputDataFolder, S3InputDataFolder
 
         if path.startswith("s3://"):
@@ -91,7 +170,18 @@ def from_path(cls, path, **kwargs):
         return LocalInputDataFolder(path, **kwargs)
 
     @abstractmethod
-    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[InputDataFile]:
+    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[BaseInputDataFile]:
+        """Retrieves a list of InputDataFile for all files in this folder matching both the dataclass properties and the
+        function parameters.
+
+        Args:
+            extension: optionally limit the file extension of the
+                retrieved files
+            suffix: optionally limit the search to a given subfolder
+
+        Returns:
+
+        """
         logger.error(
             "Do not instantiate BaseInputDataFolder directly, "
             "use a LocalInputDataFolder, S3InputDataFolder or call"
@@ -99,27 +189,81 @@ def list_files(self, extension: str | list[str] = None, suffix: str = "") -> lis
         )
         raise NotImplementedError
 
-    def __post_init__(self):
-        self._lock = contextlib.nullcontext()
-
     def set_lock(self, lock):
+        """Pass a synchronization primitive to limit concurrent downloads
+
+        Args:
+            lock: synchronization primitive (Semaphore, Lock)
+
+        Returns:
+
+        """
         self._lock = lock
 
-    def get_files_shard(self, rank: int, world_size: int, extension: str | list[str] = None) -> list[InputDataFile]:
+    def get_files_shard(
+        self, rank: int, world_size: int, extension: str | list[str] = None
+    ) -> list[BaseInputDataFile]:
+        """Fetch a shard (set of files) for a given rank, assuming there are a total of `world_size` shards.
+        This should be deterministic to not have any overlap among different ranks.
+
+        Args:
+            rank: rank of the shard to fetch
+            world_size: total number of shards
+            extension: optional file extension to pass on to list_files
+
+        Returns:
+            a list of input files
+        """
         return self.list_files(extension=extension)[rank::world_size]
 
-    def get_file(self, relative_path: str) -> InputDataFile | None:
+    def get_file(self, relative_path: str) -> BaseInputDataFile | None:
+        """Get a file directly by its name.
+
+        Args:
+            relative_path: The file name/path from this folder.
+
+        Returns:
+            an input file if it exists or None
+        """
         if self.file_exists(relative_path):
-            return self.unchecked_get_file(relative_path)
+            return self._unchecked_get_file(relative_path)
+
+    @abstractmethod
+    def _unchecked_get_file(self, relative_path: str) -> BaseInputDataFile:
+        """Get a file directly by its name, without checking if it exists.
+        Subclasses should override this method (and not the one above).
+        Should not be called directly. Instead, call `get_file`
+
+        Args:
+            relative_path: The file name/path from this folder.
 
-    def unchecked_get_file(self, relative_path: str) -> InputDataFile:
-        return InputDataFile(path=os.path.join(self.path, relative_path), relative_path=relative_path)
+        Returns:
+            an input file
+        """
+        raise NotImplementedError
 
     @abstractmethod
     def file_exists(self, relative_path: str) -> bool:
+        """Should be overriden by subclasses. Check if a given file exists.
+
+        Args:
+            relative_path: The file name/path from this folder.
+
+        Returns:
+            if the file exists
+        """
         return True
 
-    def _match_file(self, file_path, extension=None):
+    def _match_file(self, file_path, extension=None) -> bool:
+        """Checks if a given file matches the chosen extension(s) and/or pattern.
+
+        Args:
+            file_path: the relative file path
+            extension
+
+        Returns:
+            bool
+        """
         extensions = (
             ([self.extension] if isinstance(self.extension, str) else self.extension)
             if not extension
@@ -132,52 +276,198 @@ def _match_file(self, file_path, extension=None):
         )
 
 
-def get_extension(filepath):
+def get_extension(filepath, depth=None) -> str:
+    """Get full file extension (example: .jsonl.gz)
+    Optionally only get last `depth` extensions (depth=1: .gz)
+
+    Args:
+        filepath: relative path to the file
+        depth: how many extensions maximum to get
+
+    Returns:
+        the extension
+    """
     exts = []
     stem, ext = os.path.splitext(filepath)
     while ext:
         exts.append(ext)
         stem, ext = os.path.splitext(stem)
+        if depth and len(exts) >= depth:
+            break
     return "".join(reversed(exts))
 
 
+def guess_compression(filename) -> str | None:
+    """Guess compression scheme from file extension
+
+    Args:
+        filename
+
+    Returns:
+
+    """
+    match get_extension(filename, depth=1):
+        case ".gz":
+            return "gzip"
+        case ".zst":
+            return "zst"
+        case _:
+            return None
+
+
 @dataclass
-class OutputDataFile(ABC):
-    local_path: str | None
+class BaseOutputDataFile(ABC):
+    """Represents an individual output file that can be written to.
+    Supports writing directly to a gzip compressed file
+
+    Args:
+        path (str): absolute path of this file in its original form (example: s3://...)
+        relative_path (str): path relative to the output data folder this file belongs to
+    """
+
     path: str
     relative_path: str
-    file_handler = None
-    nr_documents: int = 0
+    _file_handler = None
+    _mode: str = None
+    _lock: SemLock | contextlib.nullcontext = field(default_factory=contextlib.nullcontext)
+
+    @classmethod
+    def from_path(cls, path: str, **kwargs):
+        """OutputDataFile factory: get the correct instance from a path. Additional arguments will be passed to the
+        constructor of the matched implementation.
+
+        Args:
+            path: the full path to match
+            **kwargs: any additional argument passed to the matched
+                implementation
+
+        Returns:
+
+        """
+        from datatrove.io import FSSpecOutputDataFile, LocalOutputDataFile, S3OutputDataFile
+
+        if "relative_path" not in kwargs:
+            kwargs["relative_path"] = os.path.basename(path)
+
+        if path.startswith("s3://"):
+            return S3OutputDataFile(path, **kwargs)
+        elif valid_fsspec_path(path):
+            return FSSpecOutputDataFile(path, **kwargs)
+        return LocalOutputDataFile(path, **kwargs)
 
     def close(self):
-        if self.file_handler:
-            self.file_handler.close()
+        """Close the underlying file object.
+        Subclasses my also save/upload the file
+
+        Returns:
+
+        """
+        if self._file_handler:
+            self._file_handler.close()
+            self._file_handler = None
+
+    def __enter__(self):
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def open(self, mode: str = "w", gzip: bool = False, overwrite: bool = False):
-        if not self.file_handler or overwrite:
-            os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
-            self.file_handler = open(self.local_path, mode) if not gzip else gzip_lib.open(self.local_path, mode)
+    def open(self, mode: str = "w", gzip: bool = False):
+        """Open/create the underlying file object
+
+        Args:
+            mode: mode to open
+            gzip: whether to compress the file with gzip
+
+        Returns:
+
+        """
+        if not self._file_handler or self._mode != mode:
+            if self._file_handler:
+                self._file_handler.close()
+            self._file_handler = self._create_file_handler(mode, gzip)
+            self._mode = mode
         return self
 
+    @abstractmethod
+    def _create_file_handler(self, mode: str = "w", gzip: bool = False):
+        """Should be overriden by subclasses. Opens/creates the underlying file object
+
+        Args:
+            mode
+            gzip
+
+        Returns:
+
+        """
+        raise NotImplementedError
+
     def write(self, *args, **kwargs):
-        self.file_handler.write(*args, **kwargs)
+        """Write to the underlying file object
+
+        Args:
+            *args
+            **kwargs
+
+        Returns:
+
+        """
+        self._file_handler.write(*args, **kwargs)
+
+    def get_mmap(self, dtype):
+        """Get a numpy memmap for this file
+
+        Args:
+            dtype: memmap dtype
+
+        Returns:
+
+        """
+        return np.memmap(self._file_handler, dtype=dtype)
+
+    def delete(self):
+        """Remove pointer to the given file, close its underlying file object and delete it from disk if it exists locally.
+        Does not call close() on the OutputDataFile directly to avoid uploading/saving it externally.
+
+        Returns:
+
+        """
+        if self._file_handler:
+            self._file_handler.close()
+            self._file_handler = None
 
 
 @dataclass
 class BaseOutputDataFolder(ABC):
+    """Base output data folder class. Specific implementations should override its relevant methods.
+
+    Args:
+        path (str): full path to the folder, respecting the format of specific implementations (i.e. s3 paths should start with s3://)
+    """
+
     path: str
-    local_path: str
-    _output_files: dict[str, OutputDataFile] = field(default_factory=dict)
+    _output_files: dict[str, BaseOutputDataFile] = field(default_factory=dict)
+    _lock: SemLock | contextlib.nullcontext = field(default_factory=contextlib.nullcontext)
 
     def close(self):
+        """Close all the output files and cleanup any leftover temporary files."""
         for file in self._output_files.values():
             file.close()
+        self._output_files.clear()
 
     @classmethod
     def from_path(cls, path: str, **kwargs):
+        """OutputDataFolder factory: get the correct instance from a path. Additional arguments will be passed to the
+        constructor of the matched implementation.
+
+        Args:
+            path: the full path to match
+            **kwargs: any additional argument passed to the matched
+                implementation
+
+        Returns:
+
+        """
         from datatrove.io import FSSpecOutputDataFolder, LocalOutputDataFolder, S3OutputDataFolder
 
         if path.startswith("s3://"):
@@ -187,28 +477,48 @@ def from_path(cls, path: str, **kwargs):
         return LocalOutputDataFolder(path, **kwargs)
 
     @abstractmethod
-    def create_new_file(self, relative_path: str) -> OutputDataFile:
+    def create_new_file(self, relative_path: str) -> BaseOutputDataFile:
+        """Create an OutputDataFile for the file located on `relative_path`, its path relative to this folder.
+        Each io implementation should override it.
+        Usually shouldn't be called directly - call open() instead.
+
+        Args:
+            relative_path
+
+        Returns:
+            an OutputDataFile
+        """
         logger.error(
             "Do not instantiate a BaseOutputDataFolder directly, " "use a LocalOutputDataFolder or S3OutputDataFolder"
         )
         raise NotImplementedError
 
-    def __post_init__(self):
-        self._lock = contextlib.nullcontext()
-
     def set_lock(self, lock):
+        """Pass a synchronization primitive to limit concurrent uploads
+
+        Args:
+            lock: synchronization primitive (Semaphore, Lock)
+
+        Returns:
+
+        """
         self._lock = lock
 
-    def delete_file(self, relative_path: str):
-        if relative_path in self._output_files:
-            output_file = self._output_files.pop(relative_path)
-            output_file.close()
-            if output_file.local_path and os.path.isfile(output_file.local_path):
-                os.remove(output_file.local_path)
-
-    def open(self, relative_path: str, mode: str = "w", gzip: bool = False, overwrite: bool = False):
-        if relative_path not in self._output_files or overwrite:
-            new_output_file = self.create_new_file(relative_path)
-            new_output_file.open(mode, gzip, overwrite=overwrite)
-            self._output_files[relative_path] = new_output_file
-        return self._output_files[relative_path]
+    def open(self, relative_path: str, mode: str = "w", gzip: bool = False) -> BaseOutputDataFile:
+        """Open/create output file with `relative_path` for writing.
+
+        Args:
+            relative_path: the file path relative to this folder
+            mode: the mode to open as (w, wb, etc)
+            gzip: whether to open in gzip mode
+
+        Returns:
+            OutputDataFile
+        """
+        if relative_path not in self._output_files:
+            self._output_files[relative_path] = self.create_new_file(relative_path)
+        return self._output_files[relative_path].open(mode, gzip)
+
+    @abstractmethod
+    def to_input_folder(self) -> BaseInputDataFolder:
+        raise NotImplementedError
diff --git a/src/datatrove/io/fsspec.py b/src/datatrove/io/fsspec.py
index 7160da6d..994768a7 100644
--- a/src/datatrove/io/fsspec.py
+++ b/src/datatrove/io/fsspec.py
@@ -4,39 +4,56 @@
 
 import fsspec
 
-from datatrove.io import InputDataFile
-from datatrove.io.base import BaseInputDataFolder, BaseOutputDataFolder, OutputDataFile
+from datatrove.io import BaseInputDataFile
+from datatrove.io.base import BaseInputDataFolder, BaseOutputDataFile, BaseOutputDataFolder
 
 
 @dataclass
-class FSSpecOutputDataFile(OutputDataFile):
+class FSSpecOutputDataFile(BaseOutputDataFile):
+    """An individual FSSpecOutputDataFile."""
+
     _fs: fsspec.AbstractFileSystem = None
 
-    def open(self, mode: str = "w", gzip: bool = False, overwrite: bool = False):
+    def _create_file_handler(self, mode: str = "w", gzip: bool = False):
         assert gzip is False
-        if not self.file_handler or overwrite:
-            self.file_handler = self._fs.open(self.local_path, mode)
-        return self
+        return self._fs.open(self.path, mode)
+
+    def delete(self):
+        super().delete()
+        if self._fs.isfile(self.path):
+            self._fs.rm(self.path)
 
 
 @dataclass
 class FSSpecOutputDataFolder(BaseOutputDataFolder):
+    """fsspec output data folder. Accepts any valid fsspec path schema, such as hf://datasets/...
+    Args:
+        path (str): the fsspec path to this folder
+        local_path (str): where to save output data locally before uploading
+        cleanup (str): remove local files after upload
+    """
+
     storage_options: dict = None
 
     def __post_init__(self):
-        super().__post_init__()
-        protocol, self.path = self.path.split("://")
+        if "://" in self.path:
+            protocol, self.path = self.path.split("://")
+        else:
+            protocol = "file"
         self._fs = fsspec.filesystem(protocol, **(self.storage_options if self.storage_options else {}))
 
     def create_new_file(self, relative_path: str):
         return FSSpecOutputDataFile(
-            path=os.path.join(self.path, relative_path), relative_path=relative_path, _fs=self._fs, local_path=None
+            path=os.path.join(self.path, relative_path), relative_path=relative_path, _fs=self._fs, _lock=self._lock
         )
 
+    def to_input_folder(self) -> BaseInputDataFolder:
+        return FSSpecInputDataFolder(path=self.path, storage_options=self.storage_options)
+
 
 @dataclass
-class FSSpecInputDataFile(InputDataFile):
-    _fs: fsspec.AbstractFileSystem
+class FSSpecInputDataFile(BaseInputDataFile):
+    _fs: fsspec.AbstractFileSystem = None
 
     @contextmanager
     def open_binary(self):
@@ -49,23 +66,22 @@ class FSSpecInputDataFolder(BaseInputDataFolder):
     storage_options: dict = None
 
     def __post_init__(self):
-        super().__post_init__()
         if "://" in self.path:
             protocol, self.path = self.path.split("://")
         else:
             protocol = "file"
         self._fs = fsspec.filesystem(protocol, **(self.storage_options if self.storage_options else {}))
 
-    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[InputDataFile]:
+    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[BaseInputDataFile]:
         return [
-            self.unchecked_get_file(os.path.relpath(path, self.path))
+            self._unchecked_get_file(os.path.relpath(path, self.path))
             for path in self._fs.ls(os.path.join(self.path, suffix), detail=False)
             if self._match_file(path, extension)
         ]
 
-    def unchecked_get_file(self, relative_path: str):
+    def _unchecked_get_file(self, relative_path: str):
         return FSSpecInputDataFile(
-            path=os.path.join(self.path, relative_path), relative_path=relative_path, _fs=self._fs
+            path=os.path.join(self.path, relative_path), relative_path=relative_path, _fs=self._fs, _lock=self._lock
         )
 
     def file_exists(self, relative_path: str) -> bool:
diff --git a/src/datatrove/io/local.py b/src/datatrove/io/local.py
index f242d64f..b339d46f 100644
--- a/src/datatrove/io/local.py
+++ b/src/datatrove/io/local.py
@@ -1,30 +1,72 @@
+import gzip as gzip_lib
 import os
-from dataclasses import dataclass, field
+from contextlib import contextmanager
+from dataclasses import dataclass
 
-from datatrove.io.base import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile, OutputDataFile
+from datatrove.io.base import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFile, BaseOutputDataFolder
 
 
 @dataclass
-class LocalOutputDataFolder(BaseOutputDataFolder):
-    local_path: str = field(init=False)
+class LocalOutputDataFile(BaseOutputDataFile):
+    def _create_file_handler(self, mode: str = "w", gzip: bool = False):
+        """Opens/creates the underlying file object
+
+        Args:
+            mode
+            gzip
+
+        Returns:
 
-    def __post_init__(self):
-        super().__post_init__()
-        self.local_path = self.path
+        """
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+        return open(self.path, mode) if not gzip else gzip_lib.open(self.path, mode)
 
-    def create_new_file(self, relative_path: str):
-        return OutputDataFile(
-            local_path=os.path.join(self.local_path, relative_path),
-            path=os.path.join(self.path, relative_path),
-            relative_path=relative_path,
+    def delete(self):
+        super().delete()
+        if os.path.isfile(self.path):
+            os.remove(self.path)
+
+
+@dataclass
+class LocalOutputDataFolder(BaseOutputDataFolder):
+    """Local output data folder. The output files will simply be saved to a folder on disk.
+    Args:
+        path (str): absolute path to a local folder on disk
+    """
+
+    def create_new_file(self, relative_path: str) -> LocalOutputDataFile:
+        return LocalOutputDataFile(
+            path=os.path.join(self.path, relative_path), relative_path=relative_path, _lock=self._lock
         )
 
+    def to_input_folder(self) -> BaseInputDataFolder:
+        return LocalInputDataFolder(path=self.path)
+
+
+@dataclass
+class LocalInputDataFile(BaseInputDataFile):
+    @contextmanager
+    def open_binary(self):
+        """Opens local file in binary mode
+
+        Returns:
+            binary stream
+        """
+        with open(self.path, mode="rb") as f:
+            yield f
+
 
 @dataclass
 class LocalInputDataFolder(BaseInputDataFolder):
-    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[InputDataFile]:
+    """Local input data folder. Read data already present on disk.
+    Args:
+        path (str): absolute path to a local folder on disk
+
+    """
+
+    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[LocalInputDataFile]:
         return [
-            InputDataFile(path, os.path.relpath(path, self.path))
+            self._unchecked_get_file(os.path.relpath(path, self.path))
             for path in get_local_file_list(os.path.join(self.path, suffix), self.recursive)
             if self._match_file(path, extension)
         ]
@@ -32,8 +74,31 @@ def list_files(self, extension: str | list[str] = None, suffix: str = "") -> lis
     def file_exists(self, relative_path: str) -> bool:
         return os.path.exists(os.path.join(self.path, relative_path))
 
+    def _unchecked_get_file(self, relative_path: str) -> LocalInputDataFile:
+        """Get a file directly by its name, without checking if it exists.
+
+        Args:
+            relative_path: The file name/path from this folder.
+
+        Returns:
+            an input file
+        """
+        return LocalInputDataFile(
+            path=os.path.join(self.path, relative_path), relative_path=relative_path, _lock=self._lock
+        )
+
 
 def get_local_file_list(path: str, recursive: bool = True) -> list[str]:
+    """Get a list of absolute paths to all the files in a given local folder, sorted.
+
+    Args:
+        path: The path to scan
+        recursive: if False will only scan the top level files (no
+            subdirectories)
+
+    Returns:
+        list of file paths
+    """
     filelist = []
     with os.scandir(path) as files:
         for f in files:
diff --git a/src/datatrove/io/s3.py b/src/datatrove/io/s3.py
index 6cc86018..20db3044 100644
--- a/src/datatrove/io/s3.py
+++ b/src/datatrove/io/s3.py
@@ -6,8 +6,8 @@
 
 from loguru import logger
 
-from datatrove.io import BaseOutputDataFolder, InputDataFile
-from datatrove.io.base import BaseInputDataFolder, OutputDataFile
+from datatrove.io import BaseInputDataFile, BaseOutputDataFolder
+from datatrove.io.base import BaseInputDataFolder, BaseOutputDataFile
 from datatrove.io.utils.s3 import (
     s3_download_file,
     s3_file_exists,
@@ -19,13 +19,20 @@
 
 @dataclass
 class S3OutputDataFolder(BaseOutputDataFolder):
+    """S3 output data folder. Output files will be uploaded to s3.
+    Args:
+        path (str): the s3 path to this folder
+        local_path (str): where to save output data locally before uploading
+        cleanup (str): remove local files after upload
+    """
+
     local_path: str = None
     cleanup: bool = True
 
     def __post_init__(self):
-        super().__post_init__()
         if not self.path.startswith("s3://"):
             raise ValueError("S3OutputDataFolder path must start with s3://")
+        self._tmpdir = None
         if not self.local_path:
             self._tmpdir = tempfile.TemporaryDirectory()
             self.local_path = self._tmpdir.name
@@ -40,56 +47,62 @@ def create_new_file(self, relative_path: str):
             local_path=os.path.join(self.local_path, relative_path),
             path=os.path.join(self.path, relative_path),
             relative_path=relative_path,
-            folder=self,
+            cleanup=self.cleanup,
+            _lock=self._lock,
         )
 
+    def to_input_folder(self) -> BaseInputDataFolder:
+        return S3InputDataFolder(path=self.path, local_path=self.local_path, cleanup=self.cleanup)
+
 
 @dataclass
-class S3OutputDataFile(OutputDataFile):
-    folder: "S3OutputDataFolder" = None
+class S3OutputDataFile(BaseOutputDataFile):
+    """An individual S3OutputFile. This file is uploaded to s3 when closed."""
+
+    local_path: str = None
+    cleanup: bool = True
 
     def __post_init__(self):
         if not self.path.startswith("s3://"):
             raise ValueError("S3OutputDataFile path must start with s3://")
+        assert self.local_path is not None, "S3OutputDataFile must have a local_path"
 
     def close(self):
         super().close()
-        with self.folder._lock:
-            logger.info(f'Uploading "{self.local_path}" to "{self.path}"...')
-            s3_upload_file(self.local_path, self.path)
-            logger.info(f'Uploaded "{self.local_path}" to "{self.path}".')
-        if self.folder.cleanup:
+        if self.local_path:
+            with self._lock:
+                logger.info(f'Uploading "{self.local_path}" to "{self.path}"...')
+                s3_upload_file(self.local_path, self.path)
+                logger.info(f'Uploaded "{self.local_path}" to "{self.path}".')
+            if self.cleanup:
+                os.remove(self.local_path)
+            self.local_path = None
+
+    def _create_file_handler(self, mode: str = "w", gzip: bool = False):
+        os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
+        return open(self.local_path, mode) if not gzip else gzip_lib.open(self.local_path, mode)
+
+    def delete(self):
+        super().delete()
+        if os.path.isfile(self.local_path):
             os.remove(self.local_path)
 
-    def __enter__(self):
-        self.open()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def open(self, mode: str = "w", gzip: bool = False, overwrite: bool = False):
-        if not self.local_path:
-            self._tmpdir = tempfile.TemporaryDirectory()
-            self.local_path = self._tmpdir.name
-        if not self.file_handler or overwrite:
-            os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
-            self.file_handler = open(self.local_path, mode) if not gzip else gzip_lib.open(self.local_path, mode)
-        return self
-
-    def write(self, *args, **kwargs):
-        self.file_handler.write(*args, **kwargs)
-
 
 @dataclass
-class S3InputDataFile(InputDataFile):
+class S3InputDataFile(BaseInputDataFile):
+    """An individual s3 input file.
+    Args:
+        local_path (str): local path where this file will be downloaded to (if `stream=False`)
+        stream (bool): stream the file directly from s3, without saving to disk
+    """
+
     local_path: str = None
     stream: bool = False
-    folder: "S3InputDataFolder" = None
+    cleanup: bool = True
 
     @contextmanager
     def open_binary(self):
-        if self.stream:
+        if self.stream or self.local_path is None:
             # stream
             response_stream = s3_get_file_stream(self.path)
             try:
@@ -99,7 +112,7 @@ def open_binary(self):
         else:
             # download
             if not os.path.isfile(self.local_path):
-                with self.folder._lock:
+                with self._lock:
                     logger.info(f'Downloading "{self.path}" to "{self.local_path}"...')
                     s3_download_file(self.path, self.local_path)
                     logger.info(f'Downloaded "{self.path}" to "{self.local_path}".')
@@ -107,18 +120,25 @@ def open_binary(self):
                 try:
                     yield f
                 finally:
-                    if self.folder.cleanup:
+                    if self.cleanup:
                         os.remove(self.local_path)
 
 
 @dataclass
 class S3InputDataFolder(BaseInputDataFolder):
+    """S3 input data folder
+    Args:
+        local_path (str): where to download the files to (if `stream=False`)
+        stream (bool): stream the file directly from s3, without saving to disk
+        cleanup (str): remove downloaded files from disk after they are closed
+
+    """
+
     local_path: str = None
     stream: bool = False
     cleanup: bool = True
 
     def __post_init__(self):
-        super().__post_init__()
         if not self.path.startswith("s3://"):
             raise ValueError("S3InputDataFolder path must start with s3://")
         self._tmpdir = None
@@ -126,23 +146,24 @@ def __post_init__(self):
             self._tmpdir = tempfile.TemporaryDirectory()
             self.local_path = self._tmpdir.name
 
-    def unchecked_get_file(self, relative_path: str):
+    def _unchecked_get_file(self, relative_path: str):
         return S3InputDataFile(
             path=os.path.join(self.path, relative_path),
             local_path=os.path.join(self.local_path, relative_path),
             relative_path=relative_path,
-            folder=self,
             stream=self.stream,
+            cleanup=self.cleanup,
+            _lock=self._lock,
         )
 
     def file_exists(self, relative_path: str) -> bool:
         return s3_file_exists(os.path.join(self.path, relative_path))
 
-    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[InputDataFile]:
+    def list_files(self, extension: str | list[str] = None, suffix: str = "") -> list[BaseInputDataFile]:
         return [
-            self.unchecked_get_file(os.path.join(suffix, path))
+            self._unchecked_get_file(os.path.join(suffix, path))
             for path in s3_get_file_list(
                 os.path.join(self.path, suffix), match_pattern=self.match_pattern, recursive=self.recursive
             )
-            if self._match_file(path, extension)
+            if self._match_file(os.path.join(suffix, path), extension)
         ]
diff --git a/src/datatrove/io/utils/s3.py b/src/datatrove/io/utils/s3.py
index 925a14e3..a6ed0348 100644
--- a/src/datatrove/io/utils/s3.py
+++ b/src/datatrove/io/utils/s3.py
@@ -20,10 +20,12 @@ def _get_s3_path_components(s3_path):
 
 
 def s3_file_exists(cloud_path):
-    """
-        Checks if a given file to path exists. Currently only check for s3. If path is a folder path return False
-    @params cloud_path:
-    @return: bool
+    """Checks if a given file to path exists. Currently only check for s3. If path is a folder path return False
+    Args:
+        cloud_path:
+
+    Returns:
+        bool
     """
     s3_client = boto3.client("s3")
     bucket_name, prefix = _get_s3_path_components(cloud_path)
@@ -38,9 +40,9 @@ def s3_file_exists(cloud_path):
 
 
 def s3_upload_file(local_path, cloud_path):
-    """
-    @param local_path:
-    @param cloud_path:
+    """Args:
+    local_path
+    cloud_path
     """
     if not os.path.isfile(local_path):
         raise OSError(f"File {local_path} does not exist")
@@ -52,9 +54,9 @@ def s3_upload_file(local_path, cloud_path):
 
 
 def s3_download_file(cloud_path, local_path):
-    """
-    @param cloud_path:
-    @param local_path:
+    """Args:
+    cloud_path
+    local_path
     """
     bucket_name, prefix = _get_s3_path_components(cloud_path)
     s3_resource = boto3.resource("s3")
@@ -82,12 +84,15 @@ def _match_prefix(base_prefix, prefix, match_pattern=None):
 
 
 def s3_get_file_list(cloud_path, match_pattern=None, recursive=True):
-    """
-        Get list of relative paths to files in a utils folder with a given (optional) extension
-    @param cloud_path:
-    @param match_pattern:
-    @param recursive:
-    @return:
+    """Get list of relative paths to files in a utils folder with a given (optional) extension
+
+    Args:
+        cloud_path
+        match_pattern
+        recursive
+
+    Returns:
+
     """
     s3_client = boto3.client("s3")
     bucket, main_prefix = _get_s3_path_components(cloud_path)
diff --git a/src/datatrove/pipeline/base.py b/src/datatrove/pipeline/base.py
index 233e2c17..e284e436 100644
--- a/src/datatrove/pipeline/base.py
+++ b/src/datatrove/pipeline/base.py
@@ -8,19 +8,28 @@ class PipelineStep(ABC):
     name: str = None
     type: str = None
 
-    def __init__(self, **kwargs):
-        self.stats = Stats(f"{self.__repr__()}")
+    def __init__(self):
+        self.stats = Stats(str(self))
 
-    def stat_update(self, key, value: int = 1):
-        self.stats.counter[key] += value
+    def stat_update(self, *labels, value: int = 1, unit: str = None):
+        for label in labels:
+            self.stats[label].update(value, unit)
+
+    def track_time(self, unit: str = None):
+        if unit:
+            self.stats.time_stats.unit = unit
+        return self.stats.time_stats
 
     def set_up_dl_locks(self, dl_lock, up_lock):
         pass
 
     def __repr__(self):
-        return f"{self.type} --> {self.name}"
+        return f"{self.type}: {self.name}"
 
     @abstractmethod
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         if data:
             yield from data
+
+    def __call__(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+        return self.run(data, rank, world_size)
diff --git a/src/datatrove/pipeline/dedup/bloom_filter.py b/src/datatrove/pipeline/dedup/bloom_filter.py
index e5c24f33..5bb82ced 100644
--- a/src/datatrove/pipeline/dedup/bloom_filter.py
+++ b/src/datatrove/pipeline/dedup/bloom_filter.py
@@ -45,23 +45,22 @@ def __init__(
         seed: int = 0,
         save_bloom_filter: bool = False,
         exclusion_writer: DiskWriter = None,
-        **kwargs,
     ):
+        """Args:
+        output_folder: output folder: local or on S3
+        m_bytes: bloom filter size in bytes (actual size x8 bigger)
+        k: number of hashes
+        expected_elements: expected number of elements, aka
+            shingles.
+        duplicate_threshold: above which documents are considered as
+            duplicated
+        n_grams: n_grams to use
+        seed: seed
+        save_bloom_filter: if true saves bloom filter for later use
+        exclusion_writer: saves duplicated data
         """
 
-        :param output_folder: output folder: local or on S3
-        :param m_bytes: bloom filter size in bytes (actual size x8 bigger)
-        :param k: number of hashes
-        :param expected_elements: expected number of elements, aka shingles.
-        :param duplicate_threshold: above which documents are considered as duplicated
-        :param n_grams: n_grams to use
-        :param seed: seed
-        :param save_bloom_filter: if true saves bloom filter for later use
-        :param exclusion_writer: saves duplicated data
-        :param kwargs:
-        """
-
-        super().__init__(**kwargs)
+        super().__init__()
         self.output_folder = output_folder
         self.m_bytes = m_bytes  # size in bits
         self.k = k if k else get_optimal_k(self.m, expected_elements=expected_elements)
@@ -150,17 +149,18 @@ def step(self, doc: Document) -> bool:
             return False
         return True
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         with self.exclusion_writer if self.exclusion_writer else contextlib.nullcontext() as writer:
             for doc_idx, doc in enumerate(data):
-                with self.stats.time_manager:
+                with self.track_time():
                     self.stat_update(StatHints.total)
-                    if self.step(doc):
-                        self.stat_update(StatHints.forwarded)
-                        yield doc
-                    else:
+                    if not self.step(doc):
+                        self.stat_update(StatHints.dropped)
                         if self.exclusion_writer:
                             writer.write(doc, rank)
+                        continue
+                self.stat_update(StatHints.forwarded)
+                yield doc
             if self.save_bloom_filter:
                 with self.output_folder.open("bloom_filter.bloom", mode="wb") as f:
                     f.write(self.bit_vector)
diff --git a/src/datatrove/pipeline/dedup/exact_substrings.py b/src/datatrove/pipeline/dedup/exact_substrings.py
index d7d32986..600b3742 100644
--- a/src/datatrove/pipeline/dedup/exact_substrings.py
+++ b/src/datatrove/pipeline/dedup/exact_substrings.py
@@ -1,5 +1,4 @@
-"""
-Here we implement exact substrings as suggested by https://arxiv.org/pdf/2107.06499.pdf.
+"""Here we implement exact substrings as suggested by https://arxiv.org/pdf/2107.06499.pdf.
 We use suffix array to deduplicate exact substrings above a minimum threshold. We will take the code tom build the
 actual suffix array and find duplicates form the GitHub page of "Deduplicating Training Data Makes Language Models
 Better".
@@ -11,21 +10,18 @@
  ... call deduplicate-text-datasets scripts
      in particular `cargo run self-similar ...` and `cargo run self-similar` need to be called
 
-
-
 3) DedupReader reads docs and ranges at the same time and remove duplicates.
 
-
 """
 import struct
-from typing import Generator
+from typing import Generator, Literal
 
 import numpy as np
 import tokenizers
 from loguru import logger
 from nltk.tokenize import word_tokenize
 
-from datatrove.io import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFolder
 from datatrove.pipeline.base import DocumentsPipeline, PipelineStep
 from datatrove.pipeline.readers import JsonlReader
 
@@ -43,8 +39,7 @@ def prepare_doc(tokenizer, doc: str, rank: int, doc_id: int):
 
 
 class DatasetToSequence(PipelineStep):
-    """
-    STAGE 1
+    """STAGE 1
     Creates a sequence of all docs pre-prepended by a unique separator. It also saves a second file with the
     bytes offset of where each individual doc begins.
     """
@@ -52,14 +47,12 @@ class DatasetToSequence(PipelineStep):
     type = "🫂 - DEDUP"
     name = "🪞 - exact-substrings stage 1"
 
-    def __init__(self, output_folder=BaseOutputDataFolder, tokenizer_name: str = "gpt2", **kwargs):
+    def __init__(self, output_folder=BaseOutputDataFolder, tokenizer_name: str = "gpt2"):
+        """Args:
+        output_folder: folder where sequences are saved
+        tokenizer_name: name of tokenizer as in HF tokenizers.
         """
-
-        :param output_folder: folder where sequences are saved
-        :param tokenizer_name: name of tokenizer as in HF tokenizers.
-        :param kwargs:
-        """
-        super().__init__(**kwargs)
+        super().__init__()
         self.output_folder = output_folder
         self.tokenizer = tokenizers.Tokenizer.from_pretrained(tokenizer_name)
 
@@ -68,16 +61,16 @@ def set_up_dl_locks(self, dl_lock, up_lock):
 
     def save_sizes(self, doc_lens: list[int], rank: int):
         f_lens = self.output_folder.open(f"{rank:05d}{EH.stage_1_sequence_size}", mode="wb")
-        f_lens.file_handler.write(struct.pack("Q" * len(doc_lens), *doc_lens))
+        f_lens.write(struct.pack("Q" * len(doc_lens), *doc_lens))
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         doc_lens = []
         f_sequence = self.output_folder.open(f"{rank:05d}{EH.stage_1_sequence}", mode="wb")
         for i, doc in enumerate(data):
-            with self.stats.time_manager:
+            with self.stats.time_stats:
                 b_doc = prepare_doc(tokenizer=self.tokenizer, doc=doc.content, rank=rank, doc_id=i)
                 doc_lens.append(len(b_doc))
-                f_sequence.file_handler.write(b_doc)
+                f_sequence.write(b_doc)
 
         assert i < 2**32, "doc ID overflow"
         assert i + 1 == len(doc_lens), f"{i=} but {len(doc_lens)=}"
@@ -87,8 +80,7 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
 
 
 class MergeSequences(PipelineStep):
-    """
-    STAGE 2
+    """STAGE 2
     It merges all the sequences from stage 1 into a big sequence. It saves a file with the cumulative bytes offset
     of every single sequence.
     """
@@ -102,17 +94,14 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         tasks_stage_1: int,
         bytes_per_batch: int = int(500e6),
-        **kwargs,
     ):
+        """Args:
+        input_folder: folder where sequences were saved in stage 1
+        output_folder: folder where the big sequence will be saved
+        tasks_stage_1: number of tasks used in stage 1
+        bytes_per_batch: number of bytes read per sequence
         """
-
-        :param input_folder: folder where sequences were saved in stage 1
-        :param output_folder: folder where the big sequence will be saved
-        :param tasks_stage_1: number of tasks used in stage 1
-        :param bytes_per_batch: number of bytes read per sequence
-        :param kwargs:
-        """
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.tasks_stage_1 = tasks_stage_1
@@ -121,11 +110,11 @@ def __init__(
     def set_up_dl_locks(self, dl_lock, up_lock):
         self.output_folder.set_lock(up_lock)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1):
         bytes_per_sequence = [0]
-        with self.stats.time_manager:
+        with self.stats.time_stats:
             assert world_size == 1, f"{world_size=} can't be greater than 1!"
-            all_files: list[InputDataFile] = self.input_folder.list_files(extension=EH.stage_1_sequence)
+            all_files: list[BaseInputDataFile] = self.input_folder.list_files(extension=EH.stage_1_sequence)
             assert len(all_files) == self.tasks_stage_1
             f_sequence = self.output_folder.open(f"dataset{EH.stage_2_big_sequence}", mode="wb")
             for file in all_files:
@@ -133,14 +122,14 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
                 with file.open(binary=True) as f:
                     while True:
                         sequence = f.read(self.bytes_per_batch)
-                        f_sequence.file_handler.write(sequence)
+                        f_sequence.write(sequence)
                         len_sequence += len(sequence)
                         if len(sequence) != self.bytes_per_batch:
                             break
                     bytes_per_sequence.append(bytes_per_sequence[-1] + len_sequence)
 
             f_bytes = self.output_folder.open(f"bytes_offsets{EH.stage_2_bytes_offset}", mode="wb")
-            f_bytes.file_handler.write(np.array([bytes_per_sequence], np.uint32).tobytes())
+            f_bytes.write(np.array([bytes_per_sequence], np.uint32).tobytes())
             self.output_folder.close()
 
 
@@ -149,7 +138,7 @@ def read_bytes(x):
     return np.frombuffer(x[SEPARATOR_BYTES:], dtype=np.uint16).tolist()
 
 
-def sequence_reader(file: InputDataFile, size_file: InputDataFile) -> Generator[list, None, None]:
+def sequence_reader(file: BaseInputDataFile, size_file: BaseInputDataFile) -> Generator[list, None, None]:
     with size_file.open(binary=True) as f_size:
         with file.open(binary=True) as f:
             while True:
@@ -169,12 +158,11 @@ def __init__(
         self,
         data_folder: BaseInputDataFolder,
         sequence_folder: BaseInputDataFolder,
-        gzip: bool = True,
+        compression: Literal["guess", "gzip", "zst"] | None = "guess",
         tokenizer_name: str = "gpt2",
         min_doc_words: int = 50,
-        **kwargs,
     ):
-        super().__init__(data_folder=data_folder, gzip=gzip, **kwargs)
+        super().__init__(data_folder=data_folder, compression=compression)
         self.sequence_folder = sequence_folder
         self.tokenizer = tokenizers.Tokenizer.from_pretrained(tokenizer_name)
         self.min_doc_words = min_doc_words
@@ -194,13 +182,13 @@ def reset(self):
         self.rank = None
 
     def get_sequence_bytes_offset(self):
-        offset_array_file: InputDataFile = self.sequence_folder.list_files(extension=EH.stage_2_bytes_offset)[0]
+        offset_array_file: BaseInputDataFile = self.sequence_folder.list_files(extension=EH.stage_2_bytes_offset)[0]
         with offset_array_file.open(binary=True) as f:
             offset_array = f.read()
         self.sequence_bytes_offset = np.frombuffer(offset_array, dtype=np.uint32)
         logger.info(f"{self.rank=}, -> {self.sequence_bytes_offset[self.rank]=}")
 
-    def get_bytearange(self, bytes_range_file: InputDataFile):
+    def get_bytearange(self, bytes_range_file: BaseInputDataFile):
         with bytes_range_file.open(binary=False) as f:
             dup_ranges = f.read()
 
@@ -255,8 +243,7 @@ def normalize_range(self, a, b, bytes_len):
         return a, b
 
     def get_duplicate_range(self, bytes_len: int):
-        """
-        Ranges produced by deduplicate-text-dataset can fall in one of the following 4 categories
+        """Ranges produced by deduplicate-text-dataset can fall in one of the following 4 categories
 
                    left    )  A   *    B    *       A --> *, idx <-- idx + 1
                    centre  )  *   A    B    *       idx <-- idx + 1
@@ -328,7 +315,7 @@ def remove_duplicate(self, doc, bytes_content):
 
         return True
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         self.reset()
         self.rank = rank
         # loads the sequence file from stage 1, the size file from stage 1 and the bytearange file.
@@ -338,14 +325,14 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1)
             data = self.read_files_shard(self.data_folder.get_files_shard(self.rank, world_size))
         # data is still useful for the metadata lost in the sequence format.
         for doc, doc_content in zip(data, sequence_reader(sequence_file, size_file)):
-            with self.stats.time_manager:
+            with self.stats.time_stats:
                 # We check that the two generators are synced, meaning the docs sizes bytes are correct.
                 assert doc.content == self.tokenizer.decode(
                     read_bytes(doc_content), skip_special_tokens=False
                 ), f"{doc.content}\n\n{self.tokenizer.decode(read_bytes(doc_content))}"
                 to_yield = self.remove_duplicate(doc, doc_content)
             if to_yield:
-                self.stats.doc_len.update(len(doc.content))
+                self.stats.doc_len_stats += len(doc.content)
                 yield doc
 
         # we check bytes counter matches with the offset of the following rank
diff --git a/src/datatrove/pipeline/dedup/minhash.py b/src/datatrove/pipeline/dedup/minhash.py
index 1f91a81a..1e478e64 100644
--- a/src/datatrove/pipeline/dedup/minhash.py
+++ b/src/datatrove/pipeline/dedup/minhash.py
@@ -9,7 +9,7 @@
 from nltk import ngrams, word_tokenize
 
 from datatrove.data import DocumentsPipeline
-from datatrove.io import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
 from datatrove.pipeline.dedup.utils import read_tuples_from_file, sha1_hash32, sha1_hash64, simplify_content
 from datatrove.pipeline.writers.disk_base import DiskWriter
@@ -58,24 +58,18 @@ def __str__(self):
 DEFAULT_MINHASH_CONFIG = MinhashConfig()
 
 
-@dataclass
+@dataclass(order=True)
 class HashSig:
     sig: tuple[int]
-    doc_id: int
     file_id: int
+    doc_id: int
     reader_id: int
 
-    def to_tuple(self):
-        return self.sig, self.file_id, self.doc_id, self.reader_id
-
     def is_from_index(self):
         return self.reader_id != self.file_id
 
-    def __lt__(self, other):
-        return self.to_tuple() < other.to_tuple()
 
-
-def read_sigs(file: InputDataFile, reader_id: int, config: MinhashConfig, index_file: bool = False) -> Generator:
+def read_sigs(file: BaseInputDataFile, reader_id: int, config: MinhashConfig, index_file: bool = False) -> Generator:
     if index_file:
         for data in read_tuples_from_file(file, f"{config.hashes_per_bucket}{config.hash_format}"):
             yield HashSig(sig=data, doc_id=-1, file_id=-1, reader_id=reader_id)
@@ -92,9 +86,8 @@ def __init__(
         self,
         output_folder: BaseOutputDataFolder,
         config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.output_folder = output_folder
         self.config = config
         self.num_hashes = self.config.num_buckets * self.config.hashes_per_bucket
@@ -134,28 +127,27 @@ def get_shingles(self, text):
             dtype=np.uint64,
         )
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         buckets = [
             self.output_folder.open(f"bucket_{bi:03d}/{rank:05d}.minhash.sig", mode="wb")
             for bi in range(self.config.num_buckets)
         ]
         for doc_idx, doc in enumerate(data):
-            self.stat_update(StatHints.total)
-            shingles = self.get_shingles(doc.content)
-            if shingles.size != 0:
-                sig = self.get_signature(shingles)
-                for bi, (bucket, bucket_sig) in enumerate(zip(buckets, sig)):
-                    # print(f"{self.hashes_per_bucket=} {bucket_sig=}")
-                    bucket.write(
-                        struct.pack(f"<{self.config.hashes_per_bucket}{self.config.hash_format}", *bucket_sig)
-                    )
-                    bucket.write(struct.pack("<I", doc_idx))
+            with self.track_time():
+                self.stat_update(StatHints.total)
+                shingles = self.get_shingles(doc.content)
+                if shingles.size != 0:
+                    sig = self.get_signature(shingles)
+                    for bi, (bucket, bucket_sig) in enumerate(zip(buckets, sig)):
+                        # print(f"{self.hashes_per_bucket=} {bucket_sig=}")
+                        bucket.write(
+                            struct.pack(f"<{self.config.hashes_per_bucket}{self.config.hash_format}", *bucket_sig)
+                        )
+                        bucket.write(struct.pack("<I", doc_idx))
         logger.info("Sorting buckets...")
         for bi, bucket in enumerate(buckets):
-            bucket.file_handler.close()
-            fo = self.output_folder.open(bucket.relative_path, mode="r+b", overwrite=True)
-            mmap = np.memmap(
-                fo.file_handler,
+            fo = self.output_folder.open(bucket.relative_path, mode="r+b")
+            mmap = fo.get_mmap(
                 dtype=[(str(i), self.config.hash_dtype) for i in range(self.config.hashes_per_bucket)]
                 + [(str(self.config.hashes_per_bucket), np.uint32)],
             )  # doc_id at the end
@@ -176,9 +168,8 @@ def __init__(
         index_folder: BaseInputDataFolder = None,
         config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         only_dedup_in_index: bool = True,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.index_folder = index_folder
@@ -189,7 +180,7 @@ def set_up_dl_locks(self, dl_lock, up_lock):
         self.input_folder.set_lock(dl_lock)
         self.output_folder.set_lock(up_lock)
 
-    def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline = None, bucket: int = 0, world_size: int = 1):
         assert data is None, "You should not use an input block before MinhashDedupBuckets"
         assert world_size == self.config.num_buckets, "You must run exactly one task per bucket"
         sig_files = self.input_folder.list_files(suffix=f"bucket_{bucket:03d}")
@@ -210,20 +201,23 @@ def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1
         out_f = self.output_folder.open(f"{bucket:05d}.dups", mode="wb")
 
         last: HashSig | None = None
-        while pq:
-            v: HashSig = heapq.heappop(pq)
-            if last and last.sig == v.sig and not v.is_from_index():
-                # write (file_id1, doc_id1, file_id2, doc_id2)
-                if last.is_from_index():
-                    # we can't actually write -1
-                    out_f.write(struct.pack("<4I", SENTINEL, SENTINEL, v.file_id, v.doc_id))
-                # if there isn't an index, or we are not only deduping in relation to the index
-                elif not index_files or not self.only_dedup_in_index:
-                    out_f.write(struct.pack("<4I", last.file_id, last.doc_id, v.file_id, v.doc_id))
-            last = v
-            next_sig = next(sig_readers[v.reader_id], None)
-            if next_sig:
-                heapq.heappush(pq, next_sig)
+        with self.track_time():
+            while pq:
+                v: HashSig = heapq.heappop(pq)
+                if last and last.sig == v.sig and not v.is_from_index():
+                    # write (file_id1, doc_id1, file_id2, doc_id2)
+                    if last.is_from_index():
+                        # we can't actually write -1
+                        out_f.write(struct.pack("<4I", SENTINEL, SENTINEL, v.file_id, v.doc_id))
+                        self.stat_update("index_match", "total_matches")
+                    # if there isn't an index, or we are not only deduping in relation to the index
+                    elif not index_files or not self.only_dedup_in_index:
+                        out_f.write(struct.pack("<4I", last.file_id, last.doc_id, v.file_id, v.doc_id))
+                        self.stat_update("total_matches")
+                last = v
+                next_sig = next(sig_readers[v.reader_id], None)
+                if next_sig:
+                    heapq.heappush(pq, next_sig)
         self.output_folder.close()
 
 
@@ -237,9 +231,8 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         save_cluster_id: bool = False,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.config = config
@@ -249,7 +242,7 @@ def set_up_dl_locks(self, dl_lock, up_lock):
         self.input_folder.set_lock(dl_lock)
         self.output_folder.set_lock(up_lock)
 
-    def __call__(self, data: DocumentsPipeline, _: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline = None, _: int = 0, world_size: int = 1):
         dup_files = self.input_folder.list_files(extension=".dups")
         assert len(dup_files) == self.config.num_buckets, "There should be exactly one .dups file per bucket"
         assert world_size == 1, "World size must be 1 for clustering"
@@ -261,24 +254,28 @@ def parent(x):
             union_set[x] = parent(union_set[x])
             return union_set[x]
 
-        for dup_file in dup_files:
-            for f1, d1, f2, d2 in read_tuples_from_file(dup_file, "4I"):
-                a, b = (f1, d1), (f2, d2)
-                union_set[parent(b)] = parent(a)
-
-        ci = 0
-        cluster_ids = {}
-        for node in sorted(union_set.keys()):
-            file, doc = node
-            p = parent(node)
-            if node != p:
-                self.output_folder.open(f"{file:06d}.remove", mode="wb").write(struct.pack("<I", doc))
-            if self.save_cluster_id:
-                if p not in cluster_ids:
-                    cluster_ids[p] = ci
-                    ci += 1
-                self.output_folder.open(f"{file:06d}.clusters", mode="wb").write(struct.pack("<I", doc))
-                self.output_folder.open(f"{file:06d}.clusters", mode="wb").write(struct.pack("<I", cluster_ids[p]))
+        with self.track_time():
+            for dup_file in dup_files:
+                for f1, d1, f2, d2 in read_tuples_from_file(dup_file, "4I"):
+                    a, b = (f1, d1), (f2, d2)
+                    union_set[parent(b)] = parent(a)
+
+            ci = 0
+            cluster_ids = {}
+            for node in sorted(union_set.keys()):
+                self.stat_update("duplicates")
+                file, doc = node
+                p = parent(node)
+                if node != p:
+                    self.output_folder.open(f"{file:06d}.remove", mode="wb").write(struct.pack("<I", doc))
+                    self.stat_update("to_remove")
+                if self.save_cluster_id:
+                    if p not in cluster_ids:
+                        cluster_ids[p] = ci
+                        ci += 1
+                        self.stat_update("clusters")
+                    self.output_folder.open(f"{file:06d}.clusters", mode="wb").write(struct.pack("<I", doc))
+                    self.output_folder.open(f"{file:06d}.clusters", mode="wb").write(struct.pack("<I", cluster_ids[p]))
         self.output_folder.close()
 
 
@@ -291,9 +288,8 @@ def __init__(
         input_folder: BaseInputDataFolder,
         exclusion_writer: DiskWriter = None,
         load_cluster_ids: bool = False,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.data_folder = input_folder
         self.exclusion_writer = exclusion_writer
         self.load_cluster_ids = load_cluster_ids
@@ -301,7 +297,7 @@ def __init__(
     def set_up_dl_locks(self, dl_lock, up_lock):
         self.data_folder.set_lock(dl_lock)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         remove_file = self.data_folder.get_file(f"{rank:06d}.remove")
 
         clusters_data = self.data_folder.get_files_shard(rank, world_size, extension=".clusters")
@@ -312,8 +308,7 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         if not remove_file:
             logger.warning(f"No .remove file for {rank=}.")
             for doc in data:
-                self.stat_update(StatHints.total)
-                self.stat_update(StatHints.forwarded)
+                self.stat_update(StatHints.total, StatHints.forwarded)
                 yield doc
             return
         with remove_file.open_binary() as f:
@@ -334,20 +329,21 @@ def load_clusters():
 
                 next_removal = get_next()
                 for idx, doc in enumerate(data):
-                    if self.load_cluster_ids:
-                        if next_cluster and idx == next_cluster[0]:
-                            doc.metadata["minhash_cluster"] = next_cluster[1]
-                            next_cluster = next(cluster_loader, None)
-
-                    self.stat_update(StatHints.total)
-                    if next_removal == idx:
-                        # to remove
-                        self.stat_update(StatHints.dropped)
-                        if self.exclusion_writer:
-                            exc_writer.write(doc, rank)
-                        next_removal = get_next()
-                        continue
-                    self.stat_update(StatHints.forwarded)
+                    with self.track_time():
+                        if self.load_cluster_ids:
+                            if next_cluster and idx == next_cluster[0]:
+                                doc.metadata["minhash_cluster"] = next_cluster[1]
+                                next_cluster = next(cluster_loader, None)
+
+                        self.stat_update(StatHints.total)
+                        if next_removal == idx:
+                            # to remove
+                            self.stat_update(StatHints.dropped)
+                            if self.exclusion_writer:
+                                exc_writer.write(doc, rank)
+                            next_removal = get_next()
+                            continue
+                        self.stat_update(StatHints.forwarded)
                     yield doc
 
 
@@ -361,9 +357,8 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         index_name: str,
         config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.config = config
@@ -373,8 +368,8 @@ def set_up_dl_locks(self, dl_lock, up_lock):
         self.input_folder.set_lock(dl_lock)
         self.output_folder.set_lock(up_lock)
 
-    def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1):
-        assert data is None, "You should not use an input block before MinhashDedupBuckets"
+    def run(self, data: DocumentsPipeline = None, bucket: int = 0, world_size: int = 1):
+        assert data is None, "You should not use an input block before MinhashBuildIndex"
         assert world_size == self.config.num_buckets, "You must run exactly one task per bucket"
         sig_files = self.input_folder.list_files(suffix=f"bucket_{bucket:03d}")
         sig_readers = [read_sigs(file, file_i, self.config) for file_i, file in enumerate(sig_files)]
@@ -386,12 +381,13 @@ def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1
         out_f = self.output_folder.open(f"bucket_{bucket:03d}/{self.index_name}.minhash.index", mode="wb")
 
         last: HashSig | None = None
-        while pq:
-            v: HashSig = heapq.heappop(pq)
-            if not last or last.sig != v.sig:
-                out_f.write(struct.pack(f"<%d{self.config.hash_format}" % self.config.hashes_per_bucket, *v.sig))
-            last = v
-            next_sig = next(sig_readers[v.file_id], None)
-            if next_sig:
-                heapq.heappush(pq, next_sig)
+        with self.track_time():
+            while pq:
+                v: HashSig = heapq.heappop(pq)
+                if not last or last.sig != v.sig:
+                    out_f.write(struct.pack(f"<%d{self.config.hash_format}" % self.config.hashes_per_bucket, *v.sig))
+                last = v
+                next_sig = next(sig_readers[v.file_id], None)
+                if next_sig:
+                    heapq.heappush(pq, next_sig)
         self.output_folder.close()
diff --git a/src/datatrove/pipeline/dedup/sentence_dedup.py b/src/datatrove/pipeline/dedup/sentence_dedup.py
index 51072b69..8cc75451 100644
--- a/src/datatrove/pipeline/dedup/sentence_dedup.py
+++ b/src/datatrove/pipeline/dedup/sentence_dedup.py
@@ -1,5 +1,4 @@
-"""
-'To deduplicate the data set, we discarded all but one of any three-sentence span
+"""'To deduplicate the data set, we discarded all but one of any three-sentence span
 occurring more than once in the data set.'
 
 from: https://jmlr.org/papers/volume21/20-074/20-074.pdf (C4)
@@ -19,7 +18,7 @@
 from nltk.tokenize import sent_tokenize, word_tokenize
 
 from datatrove.data import Document, DocumentsPipeline
-from datatrove.io import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
 from datatrove.utils.typeshelper import StatHints
 
@@ -27,23 +26,15 @@
 from .utils import ExtensionHelperSD, merge_docs, read_tuples_from_file, simplify_content, str_hash
 
 
-@dataclass
+@dataclass(order=True)
 class HashSig:
+    # this also determines the sorting order
+    # hash_value needs to come first as that's what we match on
+    # file_id should come after doc_id so that hashes from the index (sent_id=doc_id=-1) come up first
     hash_value: int
     doc_id: int
-    sent_id: int
     file_id: int = None
-
-    def to_tuple(self):
-        # this also determines the sorting order
-        # hash_value needs to come first as that's what we match on
-        # file_id should come after doc_id so that hashes from the index (sent_id=doc_id=-1) come up first
-        # return self.hash_value, self.sent_id, self.doc_id, self.file_id
-        return self.hash_value, self.doc_id, self.file_id, self.sent_id
-
-    # priority queue accepts anything that is sortable
-    def __lt__(self, other) -> bool:
-        return self.to_tuple() < other.to_tuple()
+    sent_id: int = None
 
     def is_from_index(self):
         return self.doc_id == self.sent_id == -1
@@ -53,14 +44,13 @@ class SentenceDedupSignature(PipelineStep):
     type = "🫂 - DEDUPS"
     name = "💥 sentence-deduplication stage 1"
 
-    def __init__(self, output_folder: BaseOutputDataFolder, n_sentences: int = 3, **kwargs):
+    def __init__(self, output_folder: BaseOutputDataFolder, n_sentences: int = 3):
+        """Args:
+        output_folder: folder where signatures are saved
+        n_sentences: n_sentences where duplicates are checked.
+        kwargs
         """
-
-        :param output_folder: folder where signatures are saved
-        :param n_sentences: n_sentences where duplicates are checked.
-        :param kwargs:
-        """
-        super().__init__(**kwargs)
+        super().__init__()
         self.output_folder = output_folder
         self.n_sentences = n_sentences
 
@@ -95,13 +85,13 @@ def get_hashes(self, doc: Document, doc_idx: int) -> list[None] | list[HashSig]:
 
         return hashes
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
-        """
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+        """Args:
+            data
+            rank
+            world_size
 
-        :param data:
-        :param rank:
-        :param world_size:
-        :return:
+        Returns:
 
         SentenceDedupSignature creates a signature for each document. Each HashSig has n hash, the doc id and the
         sentence idx. Before saving them the hashes are sorted.
@@ -109,14 +99,14 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         """
         signatures = []
         for doc_idx, doc in enumerate(data):
-            with self.stats.time_manager:
+            with self.stats.time_stats:
                 self.stat_update(StatHints.total)
                 signatures.extend(self.get_hashes(doc, doc_idx))
         self.save_hashes(rank, signatures)
         self.output_folder.close()
 
 
-def read_sigs(file: InputDataFile, file_id: int, index_file: bool = False) -> Generator[HashSig, None, None]:
+def read_sigs(file: BaseInputDataFile, file_id: int, index_file: bool = False) -> Generator[HashSig, None, None]:
     if index_file:
         # only read hashes
         for (hash,) in read_tuples_from_file(file, "Q"):
@@ -136,28 +126,27 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         index_folder: BaseInputDataFolder = None,
         only_dedup_in_index: bool = True,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.data_folder = data_folder
         self.output_folder = output_folder
         self.index_folder = index_folder
         self.only_dedup_in_index = only_dedup_in_index
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
-        """
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1):
+        """Args:
+            data
+            rank
+            world_size
 
-        :param data:
-        :param rank:
-        :param world_size:
-        :return:
+        Returns:
 
         SentenceFindDedups runs on a single worker. It reads all the signatures from the previous step and load them
         in a priority queue to check for duplicates. If a duplicate is found its document id and sentence id are saved.
         """
         assert world_size == 1, "SentenceFindDedups can only run on a single worker."
         files_with_duplicates = set()
-        with self.stats.time_manager:
+        with self.stats.time_stats:
             sig_files = self.data_folder.list_files(extension=ExtensionHelperSD.stage_1_signature)
             sig_readers = [read_sigs(file, file_i) for file_i, file in enumerate(sig_files)]
             index_files = self.index_folder.list_files() if self.index_folder else None
@@ -200,7 +189,7 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         self.output_folder.close()
 
 
-def read_duplicates(file: InputDataFile) -> Generator[tuple, None, None]:
+def read_duplicates(file: BaseInputDataFile) -> Generator[tuple, None, None]:
     yield from read_tuples_from_file(file, "I", "H")  # (doc_id, sent_id) pairs
 
 
@@ -214,15 +203,12 @@ def __init__(
         n_sentences: int = 3,
         min_doc_words: int = 50,
         exclusion_writer: DiskWriter = None,
-        **kwargs,
     ):
+        """Args:
+        data_folder: data folder to get duplicate files.
+        min_doc_words: min amount of words for each document
         """
-
-        :param data_folder: data folder to get duplicate files.
-        :param min_doc_words: min amount of words for each document
-        :param kwargs:
-        """
-        super().__init__(**kwargs)
+        super().__init__()
         self.data_folder: BaseInputDataFolder = data_folder
         self.n_sentences = n_sentences
         self.min_doc_words = min_doc_words
@@ -251,18 +237,14 @@ def remove_dup_sentences(self, doc: Document, du_lines: set = None) -> (str, str
         if in_removed_span:
             original_formatted.append("<<<\u001b[0m")
         if len(kept_sentences) < len(sentence_spans):
-            self.stat_update("removed_sentences", len(sentence_spans) - len(kept_sentences))
-        self.stat_update("original_sentences", len(sentence_spans))
+            self.stat_update("removed_sentences", value=len(sentence_spans) - len(kept_sentences))
+        self.stat_update("original_sentences", value=len(sentence_spans))
         return "".join(kept_sentences).lstrip(), "".join(original_formatted)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
-        """
-        step method for Filters.
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+        """step method for Filters.
         Drops documents that if .filter() is False
 
-        @param datapipe: input DocumentsPipeline
-        @return: DocumentsPipeline
-
         SentenceDedupFilter reads a DocumentPipeline and removes duplicated sentences found at stage 2
         """
         files = self.data_folder.get_files_shard(rank, world_size, extension=ExtensionHelperSD.stage_2_duplicates)
@@ -275,12 +257,12 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1)
         with self.exclusion_writer if self.exclusion_writer else contextlib.nullcontext() as writer:
             for idx, doc in enumerate(data):
                 self.stat_update(StatHints.total)
-                with self.stats.time_manager:
+                with self.stats.time_stats:
                     filtered_content, original_formatted = self.remove_dup_sentences(doc, du_lines=du_file.get(idx))
                 if (
                     filtered_content == doc.content or len(word_tokenize(filtered_content)) > self.min_doc_words
                 ):  # document is kept
-                    self.stats.doc_len.update(len(doc.content))
+                    self.stats.doc_len_stats += len(doc.content)
                     if not filtered_content == doc.content and writer:
                         writer.write(dataclasses.replace(doc, content=original_formatted), rank=rank)
                     doc.content = filtered_content
@@ -295,16 +277,19 @@ class SentenceDedupBuildIndex(PipelineStep):
     name = "💥 sentence-deduplication build index"
 
     def __init__(
-        self, data_folder: BaseInputDataFolder, output_folder: BaseOutputDataFolder, index_name: str, **kwargs
+        self,
+        data_folder: BaseInputDataFolder,
+        output_folder: BaseOutputDataFolder,
+        index_name: str,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.data_folder = data_folder
         self.output_folder = output_folder
         self.index_name = index_name
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1):
         assert world_size == 1, "SentenceDedupBuildIndex can only run on a single worker."
-        with self.stats.time_manager:
+        with self.stats.time_stats:
             sig_files = self.data_folder.list_files(extension=ExtensionHelperSD.stage_1_signature)
             sig_readers = [read_sigs(file, file_i) for file_i, file in enumerate(sig_files)]
 
diff --git a/src/datatrove/pipeline/dedup/utils.py b/src/datatrove/pipeline/dedup/utils.py
index c99ecf45..396f6852 100644
--- a/src/datatrove/pipeline/dedup/utils.py
+++ b/src/datatrove/pipeline/dedup/utils.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from datatrove.io import InputDataFile
+from datatrove.io import BaseInputDataFile
 
 
 class ExtensionHelperSD:
@@ -28,7 +28,14 @@ class ExtensionHelperES:
 )
 
 
-def read_tuples_from_file(file: InputDataFile, *formats):
+def read_tuples_from_file(file: BaseInputDataFile, *formats):
+    """
+        Utility to easily parse binary files. formats is a list of struct format characters.
+        yields tuples of size len(formats) with the data read
+    :param file: the file to read from
+    :param formats: list of struct format chars. Example, for 2 uint32 and 1 uint64: ['I', 'I', 'Q']
+    :return: tuples with data specified in formats
+    """
     with file.open(binary=True) as f:
         while True:
             line = []
@@ -40,6 +47,19 @@ def read_tuples_from_file(file: InputDataFile, *formats):
 
 
 def simplify_content(text: str):
+    """Performs the following operations to increase recall when looking for matches between documents:
+    - lowercase text
+    - replace all whitespace with a single " "
+    - remove all punctuation
+    - convert diacritics
+    - unicode normalize
+
+    Args:
+        text
+
+    Returns:
+        modified text
+    """
     # lower case
     text = text.lower()
     # remove consecutive spaces, newlines, tabs in the middle and in the beginning / end
diff --git a/src/datatrove/pipeline/extractors/base.py b/src/datatrove/pipeline/extractors/base.py
index 6e70da7a..1a3ed95e 100644
--- a/src/datatrove/pipeline/extractors/base.py
+++ b/src/datatrove/pipeline/extractors/base.py
@@ -9,28 +9,32 @@
 
 
 class BaseExtractor(PipelineStep):
+    """Base Extractor module. Extractors extract text from html"""
+
     type = "🛢️- EXTRAC"
 
     @abstractmethod
-    def __init__(self, timeout: float = 0.1, **kwargs):
+    def __init__(self, timeout: float = 0.1):
         """
-        Base Extractor module, converts html to text.
+        :param timeout: the timeout for extraction, per document, in seconds
         """
-        super().__init__(**kwargs)
+        super().__init__()
         self.timeout = timeout
 
     @abstractmethod
     def extract(self, content: str) -> str:
-        """
-        abstract module that actually implements the extraction, e.g. trafilatura.
-        """
+        """abstract method that actually implements the extraction, e.g. trafilatura."""
         pass
 
     def timeout_extract(self, doc: Document):
-        """
-        stops the extraction if it takes longer than timeout
-        :param doc: Documnet
-        :return:
+        """Stops the extraction if it takes longer than timeout.
+        This is the main entrypoint for this class.
+
+        Args:
+            doc: Document
+
+        Returns:
+
         """
 
         def signal_handler(signum, frame):
@@ -50,15 +54,14 @@ def signal_handler(signum, frame):
         finally:
             signal.setitimer(signal.ITIMER_REAL, 0)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
-        """ """
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         for doc in data:
             self.stat_update(StatHints.total)
-            with self.stats.time_manager:
+            with self.track_time():
                 doc.content = self.timeout_extract(doc)
             if doc.content:
                 self.stat_update(StatHints.forwarded)
-                self.stats.doc_len.update(len(doc.content))
+                self.stats.doc_len_stats += len(doc.content)
                 yield doc
             else:
                 self.stat_update(StatHints.dropped)
diff --git a/src/datatrove/pipeline/extractors/modular.py b/src/datatrove/pipeline/extractors/modular.py
index 9744e455..53ac3546 100644
--- a/src/datatrove/pipeline/extractors/modular.py
+++ b/src/datatrove/pipeline/extractors/modular.py
@@ -16,16 +16,16 @@ class ReadabilityInscriptis(BaseExtractor):
     Extracts the text from the HTML document using readability and inscriptis.
     """
 
-    def __init__(self, max_new_lines: int = 2, min_text_length=25, min_text_score=20, **kwargs):
+    def __init__(self, max_new_lines: int = 2, min_text_length=25, min_text_score=20, timeout: float = 0.1):
         """
-        :param max_new_lines:
+        :param max_new_lines: maximum number of consecutive \n to keep.
         :param min_text_length: the minimum string length of a text block. If all text blocks are shorter than
         `min_text_length`, the document is considered empty.
         :param min_text_score: `score = sqrt(block_lenth - min_text_length)`. The sum of scores of all text blocks must
         be greater than `min_text_score`.
-        :param kwargs:
+        :param timeout: the timeout for extraction, per document, in seconds
         """
-        super().__init__(**kwargs)
+        super().__init__(timeout)
         self.min_text_length = min_text_length
         self.min_text_score = min_text_score
         self.new_line_chars = "\n" * max_new_lines
diff --git a/src/datatrove/pipeline/extractors/trafilatura.py b/src/datatrove/pipeline/extractors/trafilatura.py
index 4b762eca..77327bc9 100644
--- a/src/datatrove/pipeline/extractors/trafilatura.py
+++ b/src/datatrove/pipeline/extractors/trafilatura.py
@@ -4,18 +4,24 @@
 
 
 class Trafilatura(BaseExtractor):
-    """
-    Trafilatura extractor, it uses https://trafilatura.readthedocs.io/en/latest/index.html
-    """
+    """Trafilatura extractor, it uses https://trafilatura.readthedocs.io/en/latest/index.html"""
 
     name = "⛏️ Trafilatura"
 
-    def __init__(self, favour_precision: bool = False, include_images: bool = False, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, favour_precision: bool = False, include_images: bool = False, timeout: float = 0.1, **kwargs):
+        """
+
+        :param favour_precision: prefer less text but correct extraction.
+        :param include_images: not implemented currently
+        :param timeout: the timeout for extraction, per document, in seconds
+        :param kwargs: any other option will be passed to trafilatura
+        """
+        super().__init__(timeout)
         self.favour_precision = favour_precision
         self.include_images = include_images
+        self.kwargs = kwargs
         if self.include_images:
             raise NotImplementedError
 
     def extract(self, content: str) -> str:
-        return extract(content)
+        return extract(content, favor_precision=self.favour_precision, **self.kwargs)
diff --git a/src/datatrove/pipeline/filters/base_filter.py b/src/datatrove/pipeline/filters/base_filter.py
index 405ced2d..af7075be 100644
--- a/src/datatrove/pipeline/filters/base_filter.py
+++ b/src/datatrove/pipeline/filters/base_filter.py
@@ -16,40 +16,40 @@ def get_filter_result(res):
 
 
 class BaseFilter(PipelineStep, ABC):
+    """Base module for Filters. Filters remove documents."""
+
     type = "🔻 - FILTER"
 
-    def __init__(self, exclusion_writer: DiskWriter = None, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, exclusion_writer: DiskWriter = None):
+        """
+        :param exclusion_writer: optionally pass in a writer that will save the dropped documents
+        """
+        super().__init__()
         self.exclusion_writer = exclusion_writer
 
     @abstractmethod
     def filter(self, doc: Document) -> bool | Tuple[bool, str]:
-        """
-        Filter modules' main method.
-        Returns true if a sample should be filtered.
+        """Filter modules main method.
+        Returns true if a sample should be kept, false if it should be removed.
 
-        @param doc: sample to (maybe) filter
-        @return: bool - whether the doc should be filtered
-        """
-        raise NotImplementedError
-
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
-        """
-        step method for Filters.
-        Drops documents that if .filter() is False
+        Args:
+            doc: sample to filter
 
-        @param datapipe: input DocumentsPipeline
-        @return: DocumentsPipeline
+        Returns:
+            bool - whether the doc should be kept
+            or (False, str), to drop with a specific reason
         """
+        raise NotImplementedError
 
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         with self.exclusion_writer if self.exclusion_writer else contextlib.nullcontext() as writer:
             for doc in data:
                 self.stat_update(StatHints.total)
-                with self.stats.time_manager:
+                with self.track_time():
                     filter_result, reason = get_filter_result(self.filter(doc))
                     if filter_result:
                         self.stat_update(StatHints.forwarded)
-                        self.stats.doc_len.update(len(doc.content))
+                        self.stats.doc_len_stats += len(doc.content)
                     else:
                         self.stat_update(StatHints.dropped)
                         if self.exclusion_writer:
diff --git a/src/datatrove/pipeline/filters/c4_quality_filter.py b/src/datatrove/pipeline/filters/c4_quality_filter.py
index 4ed8670a..302d495b 100644
--- a/src/datatrove/pipeline/filters/c4_quality_filter.py
+++ b/src/datatrove/pipeline/filters/c4_quality_filter.py
@@ -1,34 +1,35 @@
 import re
 
+from nltk import load
 from nltk.tokenize import sent_tokenize
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
-"""
-Applies heuristic rules from C4 https://jmlr.org/papers/volume21/20-074/20-074.pdf
+class C4QualityFilter(BaseFilter):
+    """Applies heuristic rules from C4 https://jmlr.org/papers/volume21/20-074/20-074.pdf
 
-- We only retained lines that ended in a terminal punctuation mark (! . " ?)
-- We discarded any page with fewer than 5 sentences and only retained lines that contained at least 3 words
-- We removed any page that contained any word on the “List of Dirty, Naughty, Obscene or Otherwise Bad Words”
-- We removed any line with the word Javascript.
-- We removed any page where the phrase “lorem ipsum” appeared
-- We removed any pages that contained a curly bracket
-"""
+    - We only retained lines that ended in a terminal punctuation mark (! . " ?)
+    - We discarded any page with fewer than 5 sentences and only retained lines that contained at least 3 words
+    - We removed any page that contained any word on the “List of Dirty, Naughty, Obscene or Otherwise Bad Words”
+    - We removed any line with the word Javascript.
+    - We removed any page where the phrase “lorem ipsum” appeared
+    - We removed any pages that contained a curly bracket
+    """
 
+    name = "⛰️ C4 Quality"
 
-class C4QualityFilter(BaseFilter):
-    def __init__(self, **kwargs):
-        """ """
-        super().__init__(**kwargs)
+    def __init__(self, exclusion_writer: DiskWriter = None):
+        super().__init__(exclusion_writer)
 
-        self.name = "⛰️ C4 Quality"
         self.lorem_ipsum = re.compile(r"(?i)lorem ipsum")
         self.javascript = re.compile(r"(?i)javascript")
         self.min_lines = 5
         self.min_words = 3
         self.stop_chars = (".", "'", '"', "!", "?")
+        self._tokenizer = load("tokenizers/punkt/english.pickle")
 
     def line_filter(self, line: str):
         if self.javascript.search(line):
@@ -39,10 +40,11 @@ def line_filter(self, line: str):
             return False
 
     def filter(self, doc: Document) -> bool | tuple[bool, str]:
-        """
+        """Args:
+            doc
+
+        Returns:
 
-        :param doc:
-        :return:
         """
 
         lines = sent_tokenize(doc.content)
@@ -56,5 +58,6 @@ def filter(self, doc: Document) -> bool | tuple[bool, str]:
             return False, "lorem ipsum"
 
         # TODO find a way to track skip lines to re-insert them when joining lines.
+        # sent_dedup's remove_dup_sentences may be adapted for this
         lines = [line for line in lines if self.line_filter(line)]
         doc.content = " ".join(lines)
diff --git a/src/datatrove/pipeline/filters/gopher_quality_filter.py b/src/datatrove/pipeline/filters/gopher_quality_filter.py
index 8f816d4a..47b028f5 100644
--- a/src/datatrove/pipeline/filters/gopher_quality_filter.py
+++ b/src/datatrove/pipeline/filters/gopher_quality_filter.py
@@ -5,6 +5,7 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 STOP_WORDS = ["the", "be", "to", "of", "and", "that", "have", "with"]
@@ -25,7 +26,7 @@ def __init__(
         max_non_alpha_words_ratio: float | None = 0.8,
         min_stop_words: int | None = 2,
         stop_words: list[str] | None = None,
-        **kwargs,
+        exclusion_writer: DiskWriter = None,
     ):
         """
         Filter to apply Gopher's quality heuristic rules.
@@ -42,7 +43,7 @@ def __init__(
         @param min_stop_words:
         @param stop_words:
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.min_doc_words = min_doc_words
         self.max_doc_words = max_doc_words
         self.min_avg_word_length = min_avg_word_length
diff --git a/src/datatrove/pipeline/filters/gopher_repetition_filter.py b/src/datatrove/pipeline/filters/gopher_repetition_filter.py
index 93daaffa..f9b04fea 100644
--- a/src/datatrove/pipeline/filters/gopher_repetition_filter.py
+++ b/src/datatrove/pipeline/filters/gopher_repetition_filter.py
@@ -5,6 +5,7 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 """
@@ -80,13 +81,13 @@ def __init__(
         dup_para_char_frac: float | None = 0.2,
         top_n_grams: tuple[tuple[int, float]] = ((2, 0.2), (3, 0.18), (4, 0.16)),
         dup_n_grams: tuple[tuple[int, float]] = ((5, 0.15), (6, 0.14), (7, 0.13), (8, 0.12), (9, 0.11), (10, 0.10)),
-        **kwargs,
+        exclusion_writer: DiskWriter = None,
     ):
         """
 
         @param kwargs:
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
 
         self.dup_line_frac = dup_line_frac
         self.dup_para_frac = dup_para_frac
@@ -97,7 +98,6 @@ def __init__(
         self.paragraph_exp = re.compile(r"\n{2,}")
 
     def filter(self, doc: Document) -> bool | tuple[bool, str]:
-        """ """
         text = doc.content
 
         paragraphs = self.paragraph_exp.split(text.strip())
diff --git a/src/datatrove/pipeline/filters/lambda_filter.py b/src/datatrove/pipeline/filters/lambda_filter.py
index a48b1603..923610ec 100644
--- a/src/datatrove/pipeline/filters/lambda_filter.py
+++ b/src/datatrove/pipeline/filters/lambda_filter.py
@@ -2,24 +2,26 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 class LambdaFilter(BaseFilter):
     name = "👤 Lambda"
 
-    def __init__(self, filter_function: Callable[[Document], bool], **kwargs):
+    def __init__(self, filter_function: Callable[[Document], bool], exclusion_writer: DiskWriter = None):
         """
         filters documents triggering the given filter_function with respect to a specific metadata key.
 
         @param regex_exp: regex expression
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.filter_function = filter_function
 
     def filter(self, doc: Document) -> bool:
-        """
+        """Args:
+            doc: document
 
-        :param doc: document
-        :return: is_filter
+        Returns:
+            is_filter
         """
         return self.filter_function(doc)
diff --git a/src/datatrove/pipeline/filters/language_filter.py b/src/datatrove/pipeline/filters/language_filter.py
index 398551a9..7eac72ea 100644
--- a/src/datatrove/pipeline/filters/language_filter.py
+++ b/src/datatrove/pipeline/filters/language_filter.py
@@ -5,6 +5,7 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 from datatrove.utils.assets import DOWNLOAD_PATH
 from datatrove.utils.typeshelper import Languages
 
@@ -26,7 +27,7 @@ def __init__(
         languages: tuple = (Languages.english,),
         language_threshold: float = 0.65,
         model_local_path: str = os.path.join(DOWNLOAD_PATH, "language_id/lid.176.bin"),
-        **kwargs,
+        exclusion_writer: DiskWriter = None,
     ):
         """
         filters if the predicted language is not among given language or if the language score is below language
@@ -35,7 +36,7 @@ def __init__(
         @param regex_exp: language_threshold minimum score to accept a document.
         @param languages: list of languages to not filter out.
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.language_threshold = language_threshold
         self.languages = languages
         self.model_local_path = model_local_path
@@ -55,10 +56,11 @@ def model(self):
         return self._model
 
     def filter(self, doc: Document) -> bool:
-        """
+        """Args:
+            doc: document
 
-        :param doc: document
-        :return: is_filter
+        Returns:
+            is_filter
         """
 
         language, score = self.model.predict(doc.content.replace("\n", ""))
diff --git a/src/datatrove/pipeline/filters/list_filter.py b/src/datatrove/pipeline/filters/list_filter.py
index 0f206063..e967f050 100644
--- a/src/datatrove/pipeline/filters/list_filter.py
+++ b/src/datatrove/pipeline/filters/list_filter.py
@@ -2,27 +2,24 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 class ListFilter(BaseFilter):
     name = "🎅 List"
 
-    def __init__(
-        self,
-        new_line_ratio: float | None = 0.3,  # TODO better tune
-        **kwargs,
-    ):
+    def __init__(self, new_line_ratio: float | None = 0.3, exclusion_writer: DiskWriter = None):  # TODO better tune
         """ """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.new_line_ratio = new_line_ratio
 
     def filter(self, doc: Document) -> bool | tuple[bool, str]:
-        """
-            Applies the heuristics rules to decide if a document should be REMOVED:
-                -
+        """Applies heuristic rules to decide if a document should be REMOVED
+        Args:
+            doc
 
-        :param doc
-        :return: False if sample.content is a list
+        Returns:
+            False if sample.content is a list
         """
         text = doc.content
         words = word_tokenize(text)  # TODO we should use language id filter
diff --git a/src/datatrove/pipeline/filters/regex_filter.py b/src/datatrove/pipeline/filters/regex_filter.py
index 8e9e03ac..2a141ce1 100644
--- a/src/datatrove/pipeline/filters/regex_filter.py
+++ b/src/datatrove/pipeline/filters/regex_filter.py
@@ -2,24 +2,26 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 class RegexFilter(BaseFilter):
     name = "🕵️Regex"
 
-    def __init__(self, regex_exp: str, **kwargs):
+    def __init__(self, regex_exp: str, exclusion_writer: DiskWriter = None):
         """
         filters if regex find at least one match
 
         @param regex_exp: regex expression
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.regex = re.compile(regex_exp)
 
     def filter(self, doc: Document) -> bool:
-        """
+        """Args:
+            doc: document
 
-        :param doc: document
-        :return: is_filter
+        Returns:
+            is_filter
         """
         return not len(self.regex.findall(doc.content)) > 0
diff --git a/src/datatrove/pipeline/filters/sampler_filter.py b/src/datatrove/pipeline/filters/sampler_filter.py
index 683ab08e..54752713 100644
--- a/src/datatrove/pipeline/filters/sampler_filter.py
+++ b/src/datatrove/pipeline/filters/sampler_filter.py
@@ -2,19 +2,17 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
 class SamplerFilter(BaseFilter):
     name = "🎲 Sampler"
 
     def __init__(
-        self,
-        rate: float | None = 0.5,  # rate to KEEP
-        seed: int = None,
-        **kwargs,
+        self, rate: float | None = 0.5, seed: int = None, exclusion_writer: DiskWriter = None  # rate to KEEP
     ):
         """ """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.rate = rate
         self.uniform = default_rng(seed).uniform
 
diff --git a/src/datatrove/pipeline/filters/unigram_log_probs.py b/src/datatrove/pipeline/filters/unigram_log_probs.py
index fd8707b8..5537f026 100644
--- a/src/datatrove/pipeline/filters/unigram_log_probs.py
+++ b/src/datatrove/pipeline/filters/unigram_log_probs.py
@@ -7,6 +7,7 @@
 
 from datatrove.data import Document
 from datatrove.pipeline.filters.base_filter import BaseFilter
+from datatrove.pipeline.writers.disk_base import DiskWriter
 from datatrove.utils.typeshelper import LocalPaths
 
 
@@ -26,7 +27,7 @@ def __init__(
         self,
         logprobs_threshold: float = -10,
         model_local_path: str = os.path.join(LocalPaths.download, "logprob_filter/"),
-        **kwargs,
+        exclusion_writer: DiskWriter = None,
     ):
         """
         filters if the predicted language is not among given language or if the language score is below language
@@ -34,7 +35,7 @@ def __init__(
 
         @param languages: list of languages to not filter out.
         """
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.logprobs_threshold = logprobs_threshold
         self.model_local_path = model_local_path
         self.unigram_frequencies = self.get_frequencies()
diff --git a/src/datatrove/pipeline/filters/url_filter.py b/src/datatrove/pipeline/filters/url_filter.py
index a3cfa503..6e300cef 100644
--- a/src/datatrove/pipeline/filters/url_filter.py
+++ b/src/datatrove/pipeline/filters/url_filter.py
@@ -7,12 +7,10 @@
 from datatrove.data import Document
 from datatrove.utils.assets import ASSETS_PATH, DOWNLOAD_PATH
 
+from ..writers.disk_base import DiskWriter
 from .base_filter import BaseFilter
 
 
-ADULT_LIST = "ftp://ftp.ut-capitole.fr/pub/reseau/cache/squidguard_contrib/adult.tar.gz"
-
-
 normalizer = re.compile(r"[^a-zA-Z0-9]+")
 
 
@@ -36,9 +34,9 @@ def __init__(
         banned_words: set = None,
         banned_subwords: set = None,
         soft_banned_words: set = None,
-        **kwargs,
+        exclusion_writer: DiskWriter = None,
     ):
-        super().__init__(**kwargs)
+        super().__init__(exclusion_writer)
         self.download_path = os.path.join(DOWNLOAD_PATH, "url_filter")
         self.download_data()
         self.block_listed_domains = self.get_list(self.download_path, "adult/domains", extra_domains)
diff --git a/src/datatrove/pipeline/readers/base.py b/src/datatrove/pipeline/readers/base.py
index ab9f10b4..34238a16 100644
--- a/src/datatrove/pipeline/readers/base.py
+++ b/src/datatrove/pipeline/readers/base.py
@@ -1,25 +1,57 @@
 from abc import abstractmethod
 from contextlib import nullcontext
+from typing import Callable
 
 from loguru import logger
 from tqdm import tqdm
 
-from datatrove.data import DocumentsPipeline
-from datatrove.io import BaseInputDataFolder, InputDataFile
+from datatrove.data import Document, DocumentsPipeline
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 from datatrove.pipeline.base import PipelineStep
 
 
 class BaseReader(PipelineStep):
     type = "📖 - READER"
 
-    def __init__(self, data_folder: BaseInputDataFolder, limit: int = -1, progress: bool = False, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        data_folder: BaseInputDataFolder,
+        limit: int = -1,
+        progress: bool = False,
+        adapter: Callable = None,
+        content_key: str = "content",
+        id_key: str = "data_id",
+    ):
+        super().__init__()
         self.data_folder = data_folder
         self.limit = limit
         self.progress = progress
+        self.content_key = content_key
+        self.id_key = id_key
+        self.adapter = adapter if adapter else self._default_adapter
+        self.empty_warning = False
+
+    def _default_adapter(self, data: dict, path: str, id_in_file: int):
+        return {
+            "content": data.pop(self.content_key, ""),
+            "data_id": data.pop(self.id_key, f"{path}/{id_in_file}"),
+            "media": data.pop("media", []),
+            "metadata": data,
+        }
+
+    def get_document_from_dict(self, data: dict, source_file: BaseInputDataFile, id_in_file: int):
+        parsed_data = self.adapter(data, source_file.relative_path, id_in_file)
+        if not parsed_data.get("content", None):
+            if not self.empty_warning:
+                self.empty_warning = True
+                logger.warning("Found document without content, skipping.")
+            return None
+        document = Document(**parsed_data)
+        document.metadata.setdefault("file_path", source_file.path)
+        return document
 
     @abstractmethod
-    def read_file(self, datafile: InputDataFile):
+    def read_file(self, datafile: BaseInputDataFile):
         raise NotImplementedError
 
     def set_up_dl_locks(self, dl_lock, up_lock):
@@ -29,16 +61,20 @@ def read_files_shard(self, shard):
         li = 0
         with tqdm(total=self.limit if self.limit != -1 else None) if self.progress else nullcontext() as pbar:
             for datafile in shard:
+                self.stat_update("input_files")
                 logger.info(f"Reading input file {datafile.path}")
-                for document in self.read_file(datafile):
+                for di, document in enumerate(self.read_file(datafile)):
                     if self.limit != -1 and li >= self.limit:
-                        return
+                        break
                     yield document
                     if self.progress:
                         pbar.update()
                     li += 1
+                self.stat_update("documents", value=di, unit="input_file")
+                if self.limit != -1 and li >= self.limit:
+                    break
 
-    def __call__(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         if data:
             yield from data
         yield from self.read_files_shard(self.data_folder.get_files_shard(rank, world_size))
diff --git a/src/datatrove/pipeline/readers/csv.py b/src/datatrove/pipeline/readers/csv.py
index ce7b80d4..5ad344c1 100644
--- a/src/datatrove/pipeline/readers/csv.py
+++ b/src/datatrove/pipeline/readers/csv.py
@@ -1,10 +1,7 @@
 import csv
 from typing import Callable, Literal
 
-from loguru import logger
-
-from datatrove.data import Document
-from datatrove.io import BaseInputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 from datatrove.pipeline.readers.base import BaseReader
 
 
@@ -14,32 +11,23 @@ class CSVReader(BaseReader):
     def __init__(
         self,
         data_folder: BaseInputDataFolder,
-        content_column: str = "content",
-        id_column: str = "data_id",
-        compression: Literal["gzip", "zst"] | None = None,
+        compression: Literal["guess", "gzip", "zst"] | None = "guess",
+        limit: int = -1,
+        progress: bool = False,
         adapter: Callable = None,
-        **kwargs,
+        content_key: str = "content",
+        id_key: str = "data_id",
     ):
-        super().__init__(data_folder, **kwargs)
-        self.content_column = content_column
-        self.id_column = id_column
+        super().__init__(data_folder, limit, progress, adapter, content_key, id_key)
         self.compression = compression
-        self.adapter = adapter if adapter else self._base_adapter
         self.empty_warning = False
 
-    def read_file(self, datafile: InputDataFile):
+    def read_file(self, datafile: BaseInputDataFile):
         with datafile.open(compression=self.compression) as f:
             csv_reader = csv.DictReader(f)
             for di, d in enumerate(csv_reader):
-                with self.stats.time_manager:
-                    if not d.get(self.content_column, None):
-                        if not self.empty_warning:
-                            self.empty_warning = True
-                            logger.warning("Found document without content, skipping.")
+                with self.track_time():
+                    document = self.get_document_from_dict(d, datafile, di)
+                    if not document:
                         continue
-                    document = Document(**self.adapter(d, datafile.path, di))
-                    document.metadata.setdefault("file_path", datafile.path)
                 yield document
-
-    def _base_adapter(self, d: dict, path: str, di: int):
-        return {"content": d.pop(self.content_column), "data_id": d.pop(self.id_column, f"{path}_{di}"), "metadata": d}
diff --git a/src/datatrove/pipeline/readers/jsonl.py b/src/datatrove/pipeline/readers/jsonl.py
index 17bca9f2..85d22052 100644
--- a/src/datatrove/pipeline/readers/jsonl.py
+++ b/src/datatrove/pipeline/readers/jsonl.py
@@ -4,8 +4,7 @@
 
 from loguru import logger
 
-from datatrove.data import Document
-from datatrove.io import BaseInputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 from datatrove.pipeline.readers.base import BaseReader
 
 
@@ -15,30 +14,24 @@ class JsonlReader(BaseReader):
     def __init__(
         self,
         data_folder: BaseInputDataFolder,
-        compression: Literal["gzip", "zst"] | None = None,
+        compression: Literal["guess", "gzip", "zst"] | None = "guess",
+        limit: int = -1,
+        progress: bool = False,
         adapter: Callable = None,
         content_key: str = "content",
-        **kwargs,
+        id_key: str = "data_id",
     ):
-        super().__init__(data_folder, **kwargs)
+        super().__init__(data_folder, limit, progress, adapter, content_key, id_key)
         self.compression = compression
-        self.content_key = content_key
-        self.adapter = adapter if adapter else lambda d, path, li: d
-        self.empty_warning = False
 
-    def read_file(self, datafile: InputDataFile):
+    def read_file(self, datafile: BaseInputDataFile):
         with datafile.open(compression=self.compression) as f:
             for li, line in enumerate(f):
-                with self.stats.time_manager:
+                with self.track_time():
                     try:
-                        d = json.loads(line)
-                        if not d.get(self.content_key, None):
-                            if not self.empty_warning:
-                                self.empty_warning = True
-                                logger.warning("Found document without content, skipping.")
+                        document = self.get_document_from_dict(json.loads(line), datafile, li)
+                        if not document:
                             continue
-                        document = Document(**self.adapter(d, datafile.path, li))
-                        document.metadata.setdefault("file_path", datafile.path)
                     except (EOFError, JSONDecodeError) as e:
                         logger.warning(f"Error when reading `{datafile.path}`: {e}")
                         continue
diff --git a/src/datatrove/pipeline/readers/parquet.py b/src/datatrove/pipeline/readers/parquet.py
index 4baa975b..73952ab9 100644
--- a/src/datatrove/pipeline/readers/parquet.py
+++ b/src/datatrove/pipeline/readers/parquet.py
@@ -2,8 +2,7 @@
 
 import pyarrow.parquet as pq
 
-from datatrove.data import Document
-from datatrove.io import BaseInputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 from datatrove.pipeline.readers.base import BaseReader
 
 
@@ -13,17 +12,20 @@ class ParquetReader(BaseReader):
     def __init__(
         self,
         data_folder: BaseInputDataFolder,
+        limit: int = -1,
+        progress: bool = False,
         adapter: Callable = None,
-        **kwargs,
+        content_key: str = "content",
+        id_key: str = "data_id",
     ):
-        super().__init__(data_folder, **kwargs)
-        self.adapter = adapter if adapter else lambda d, path, li: d
+        super().__init__(data_folder, limit, progress, adapter, content_key, id_key)
 
-    def read_file(self, datafile: InputDataFile):
+    def read_file(self, datafile: BaseInputDataFile):
         with datafile.open(binary=True) as f:
             with pq.ParquetFile(f) as pqf:
                 for li, line in enumerate(pqf.iter_batches(batch_size=1)):
-                    with self.stats.time_manager:
-                        document = Document(**self.adapter(line.to_pydict(), datafile.path, li))
-                        document.metadata.setdefault("file_path", datafile.path)
+                    with self.track_time():
+                        document = self.get_document_from_dict(line.to_pydict(), datafile, li)
+                        if not document:
+                            continue
                     yield document
diff --git a/src/datatrove/pipeline/readers/warc.py b/src/datatrove/pipeline/readers/warc.py
index 44072c08..4eb02b8d 100644
--- a/src/datatrove/pipeline/readers/warc.py
+++ b/src/datatrove/pipeline/readers/warc.py
@@ -1,32 +1,44 @@
-from typing import Literal
+from typing import Callable, Literal
 
 import cchardet
 import magic
 from warcio.archiveiterator import ArchiveIterator
 from warcio.recordloader import ArcWarcRecord
 
-from datatrove.data import Document
-from datatrove.io import InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 from datatrove.pipeline.readers.base import BaseReader
 
 
 class WarcReader(BaseReader):
     name = "🕷️ Warc"
 
-    def __init__(self, *args, compression: Literal["gzip", "zst"] | None = None, **kwargs):
+    def __init__(
+        self,
+        data_folder: BaseInputDataFolder,
+        compression: Literal["guess", "gzip", "zst"] | None = "guess",
+        limit: int = -1,
+        progress: bool = False,
+        adapter: Callable = None,
+        content_key: str = "content",
+        id_key: str = "data_id",
+    ):
         self.compression = compression
-        super().__init__(*args, **kwargs)
+        super().__init__(data_folder, limit, progress, adapter, content_key, id_key)
 
-    def read_file(self, datafile: InputDataFile):
+    def read_file(self, datafile: BaseInputDataFile):
         with datafile.open(compression=self.compression, binary=True) as f:
-            for record in ArchiveIterator(f):
-                document = process_record(record)
-                if document:
-                    document.metadata["file_path"] = datafile.path
-                    yield document
+            for ri, record in enumerate(ArchiveIterator(f)):
+                with self.track_time():
+                    extracted_data = process_record(record)
+                    if not extracted_data:
+                        continue
+                    document = self.get_document_from_dict(extracted_data, datafile, ri)
+                    if not document:
+                        continue
+                yield document
 
 
-def process_record(record: ArcWarcRecord):
+def process_record(record: ArcWarcRecord) -> dict | None:
     # record type
     if record.rec_type != "response":
         return
@@ -67,8 +79,4 @@ def process_record(record: ArcWarcRecord):
     if not date:
         date = dict(record.rec_headers.headers)["archive-date"]
 
-    return Document(
-        content=html,
-        data_id=data_id,
-        metadata={"url": url, "date": date},
-    )
+    return {"content": html, "data_id": data_id, "url": url, "date": date}
diff --git a/src/datatrove/pipeline/stats/stats.py b/src/datatrove/pipeline/stats/stats.py
index 7f5e042a..c7e10446 100644
--- a/src/datatrove/pipeline/stats/stats.py
+++ b/src/datatrove/pipeline/stats/stats.py
@@ -5,11 +5,7 @@ class Nerdino(PipelineStep):
     type = "📊 - STATS"
     name = "🤓 document length"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         for doc in data:
-            with self.stats.time_manager:
-                self.stats.doc_len.update(len(doc.content))
+            self.stats.doc_len_stats += len(doc.content)
             yield doc
diff --git a/src/datatrove/pipeline/tokens/context_shuffler.py b/src/datatrove/pipeline/tokens/context_shuffler.py
index ddba27ba..d7dddec9 100644
--- a/src/datatrove/pipeline/tokens/context_shuffler.py
+++ b/src/datatrove/pipeline/tokens/context_shuffler.py
@@ -20,9 +20,8 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         window_size: int = 2048 + 1,
         seed: int = None,
-        **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.window_size = window_size
@@ -36,7 +35,7 @@ def get_ordering(self, all_doc_ends):
         doc_ids = np.concatenate([np.ones(len(doc_ends), dtype=int) * i for i, doc_ends in enumerate(all_doc_ends)])
         return doc_ids if not self.shuffle else self.rand.permutation(doc_ids)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         datafiles = self.input_folder.get_files_shard(rank, world_size, extension=".ds")
         datafiles_index = self.input_folder.get_files_shard(rank, world_size, extension=".ds.index")
         for datafile, index in zip(datafiles, datafiles_index):
@@ -47,7 +46,8 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1)
             fout = self.output_folder.open(datafile.relative_path, "wb")
             with datafile.open_binary() as f:
                 with mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ) as unshuf:
-                    for windowi in ordering:
-                        start, end = windowi * self.window_size * 2, (windowi + 1) * self.window_size * 2
-                        fout.write(unshuf[start:end])
+                    with self.track_time():
+                        for windowi in ordering:
+                            start, end = windowi * self.window_size * 2, (windowi + 1) * self.window_size * 2
+                            fout.write(unshuf[start:end])
         self.output_folder.close()
diff --git a/src/datatrove/pipeline/tokens/merger.py b/src/datatrove/pipeline/tokens/merger.py
index 7e181c38..d145cad1 100644
--- a/src/datatrove/pipeline/tokens/merger.py
+++ b/src/datatrove/pipeline/tokens/merger.py
@@ -4,7 +4,7 @@
 from numpy.random import default_rng
 
 from datatrove.data import DocumentsPipeline
-from datatrove.io import BaseInputDataFolder, BaseOutputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder, BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
 from datatrove.pipeline.tokens.tokenizer import TokenizedFile
 
@@ -22,10 +22,10 @@ def __init__(
         max_tokens: int = -1,  # max number of tokens to process
         shuffle: bool = True,  # whether to shuffle documents in the dataset
         seed: int = None,
-        save_loss_metadata: bool = True,
-        **kwargs,
+        save_loss_metadata: bool = False,
+        save_final_metadata: bool = True,
     ):
-        super().__init__(**kwargs)
+        super().__init__()
         self.input_folder = input_folder
         self.output_folder = output_folder
         self.save_filename = save_filename
@@ -34,6 +34,7 @@ def __init__(
         self.shuffle = shuffle
         self.save_loss_metadata = save_loss_metadata
         self.rand = default_rng(seed)
+        self.save_final_metadata = save_final_metadata
 
     def set_up_dl_locks(self, dl_lock, up_lock):
         self.input_folder.set_lock(dl_lock)
@@ -43,7 +44,7 @@ def get_ordering(self, all_doc_ends):
         doc_ids = np.concatenate([np.ones(len(doc_ends), dtype=int) * i for i, doc_ends in enumerate(all_doc_ends)])
         return doc_ids if not self.shuffle else self.rand.permutation(doc_ids)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         assert world_size == 1, "world_size must be 1 for DocumentTokenizerMerger"
         datafiles = self.input_folder.list_files(extension=".ds")
         datafiles_index = self.input_folder.list_files(extension=".ds.index")
@@ -63,6 +64,14 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1)
             if self.save_loss_metadata
             else None
         )
+
+        tokenizer_name = None
+        if self.save_final_metadata:
+            first_metadata_file = self.input_folder.get_file(f"{datafiles[0].relative_path}.metadata")
+            if first_metadata_file:
+                with first_metadata_file.open() as f:
+                    tokenizer_name = f.read().splitlines()[0]
+
         ordering = self.get_ordering(doc_ends)
 
         file_ct = 0
@@ -71,43 +80,44 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1)
             filename=f"{file_ct:03d}_{self.save_filename}.ds",
             save_loss_metadata=self.save_loss_metadata,
         )
-        output_file.open()
         for input_file_id in ordering:
-            if 0 < self.max_tokens <= len(output_file):
+            if 0 < self.max_tokens <= self.stats["tokens"].total:
                 break
             if 0 < self.max_tokens_per_file <= len(output_file):
                 output_file.close()
                 file_ct += 1
+                if self.save_final_metadata:
+                    output_file.save_final_metadata(tokenizer_name)
                 output_file = TokenizedFile(
                     output_folder=self.output_folder,
                     filename=f"{file_ct:03d}_{self.save_filename}.ds",
                     save_loss_metadata=self.save_loss_metadata,
                 )
-                output_file.open()
             # copy tokens and loss
             tokens = next(token_inputs[input_file_id])
             output_file.write_bytes(tokens)
             if loss_inputs:
                 output_file.write_loss_bytes(next(loss_inputs[input_file_id]))
-            self.stat_update("tokens", len(tokens) // 2)
+            self.stat_update("tokens", value=len(tokens) // 2)
         # cleanup
         output_file.close()
+        if self.save_final_metadata:
+            output_file.save_final_metadata(tokenizer_name)
+            # save final total metadata file
+            output_file.save_final_metadata(
+                tokenizer_name, self.stats["tokens"].total, filename=f"{self.save_filename}.ds"
+            )
         self.output_folder.close()
 
 
-def load_doc_ends(file: InputDataFile):
+def load_doc_ends(file: BaseInputDataFile):
     with file.open_binary() as f:
         return np.frombuffer(f.read(), dtype=np.uint64).tolist()
 
 
-def get_data_reader(file: InputDataFile, doc_ends: list, nb_bytes: int):
+def get_data_reader(file: BaseInputDataFile, doc_ends: list, nb_bytes: int):
     with file.open_binary() as f:
         start_e = 0
         for r_e in doc_ends:
             yield f.read((r_e - start_e) * nb_bytes)
             start_e = r_e
-
-
-# def load_input_mmap(file: InputDataFile):
-#     with file.open_binary() as f:
-#         return mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
diff --git a/src/datatrove/pipeline/tokens/tokenizer.py b/src/datatrove/pipeline/tokens/tokenizer.py
index e5141d5d..3f03dc9e 100644
--- a/src/datatrove/pipeline/tokens/tokenizer.py
+++ b/src/datatrove/pipeline/tokens/tokenizer.py
@@ -1,12 +1,13 @@
 import mmap
 import struct
 
+import humanize
 import numpy as np
 from loguru import logger
 from numpy.random import default_rng
 
 from datatrove.data import Document, DocumentsPipeline
-from datatrove.io import BaseOutputDataFolder, OutputDataFile
+from datatrove.io import BaseOutputDataFile, BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
 
 
@@ -24,28 +25,22 @@ def __init__(
         output_folder: BaseOutputDataFolder,
         filename: str,
         save_index: bool = True,
-        save_loss_metadata: bool = True,
+        save_loss_metadata: bool = False,
     ):
         self.output_folder = output_folder
         self.filename = filename
         self.save_index = save_index
         self.save_loss_metadata = save_loss_metadata
-        self.tokens_file: OutputDataFile | None = None
-        self.loss_file: OutputDataFile | None = None
         self.write_idx = 0
         self.doc_ends = []
 
-    def __len__(self):
-        return self.doc_ends[-1] if self.doc_ends else 0
-
-    def open(self):
         self.tokens_file = self.output_folder.open(self.filename, mode="wb")
+        self.loss_file: BaseOutputDataFile | None = None
         if self.save_loss_metadata:
             self.loss_file = self.output_folder.open(f"{self.filename}.loss", mode="wb")
 
-    def __enter__(self):
-        self.open()
-        return self
+    def __len__(self):
+        return self.doc_ends[-1] if self.doc_ends else 0
 
     def close(self):
         if self.tokens_file:
@@ -58,20 +53,17 @@ def close(self):
             # save total number of documents
             # index_file.file_handler.write(struct.pack('<I', len(self.doc_ends)))
             # save document boundaries - uint64
-            index_file.file_handler.write(struct.pack("<%sQ" % len(self.doc_ends), *self.doc_ends))
+            index_file.write(struct.pack("<%sQ" % len(self.doc_ends), *self.doc_ends))
             index_file.close()
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
     def cleanup(self):
         self.doc_ends = []
-        self.output_folder.delete_file(self.filename)
-        self.output_folder.delete_file(f"{self.filename}.index")
-        self.output_folder.delete_file(f"{self.filename}.loss")
+        self.tokens_file.delete()
+        if self.loss_file:
+            self.loss_file.delete()
 
     def write_bytes(self, tk_bytes: bytes):
-        self.tokens_file.file_handler.write(tk_bytes)
+        self.tokens_file.write(tk_bytes)
         # 1 token = 2 bytes (uint16)
         self.write_idx += len(tk_bytes) // 2
         # save each document's boundary
@@ -79,7 +71,7 @@ def write_bytes(self, tk_bytes: bytes):
 
     def write_loss_bytes(self, l_bytes: bytes):
         if self.save_loss_metadata:
-            self.loss_file.file_handler.write(l_bytes)
+            self.loss_file.write(l_bytes)
 
     def write(self, tokens: list[int], loss_values: np.ndarray | None):
         # get the bytes for uint16 (H)
@@ -89,32 +81,37 @@ def write(self, tokens: list[int], loss_values: np.ndarray | None):
 
     def copy(self, destination: str, ordering: np.ndarray = None):
         # open original file in read mode
-        tokens_file = self.output_folder.open(self.filename, mode="r+b", overwrite=True)
-        loss_file = (
-            None
-            if not self.loss_file
-            else self.output_folder.open(f"{self.filename}.loss", mode="r+b", overwrite=True)
-        )
-        with TokenizedFile(self.output_folder, destination, save_loss_metadata=self.save_loss_metadata) as new_file:
-            # mmap the original file
-            orig_tokens = mmap.mmap(tokens_file.file_handler.fileno(), 0)
-            orig_loss = mmap.mmap(loss_file.file_handler.fileno(), 0) if loss_file else None
-            # shuffle doc_id
-            for doc_id in ordering:
-                # get start and end from the boundaries
-                start, end = self.doc_ends[doc_id - 1] if doc_id > 0 else 0, self.doc_ends[doc_id]
-                # copy the bytes. each token is 2 bytes
-                new_file.write_bytes(orig_tokens[start * 2 : end * 2])
-                # copy loss values (1 byte per token)
+        tokens_file = self.output_folder.open(self.filename, mode="r+b")
+        loss_file = None if not self.loss_file else self.output_folder.open(f"{self.filename}.loss", mode="r+b")
+        new_file = TokenizedFile(self.output_folder, destination, save_loss_metadata=self.save_loss_metadata)
+        # mmap the original file
+        orig_tokens = mmap.mmap(tokens_file._file_handler.fileno(), 0)
+        orig_loss = mmap.mmap(loss_file._file_handler.fileno(), 0) if loss_file else None
+        # shuffle doc_id
+        for doc_id in ordering:
+            # get start and end from the boundaries
+            start, end = self.doc_ends[doc_id - 1] if doc_id > 0 else 0, self.doc_ends[doc_id]
+            # copy the bytes. each token is 2 bytes
+            new_file.write_bytes(orig_tokens[start * 2 : end * 2])
+            # copy loss values (1 byte per token)
+            if orig_loss:
                 new_file.write_loss_bytes(orig_loss[start:end])
-            # close mmaps
-            orig_tokens.close()
+        # close mmaps
+        orig_tokens.close()
+        if orig_loss:
             orig_loss.close()
-        # close files
-        tokens_file.close()
-        loss_file.close()
         return new_file
 
+    def save_final_metadata(self, tokenizer_name: str | None = None, token_count: int = -1, filename: str = None):
+        if not tokenizer_name:
+            tokenizer_name = "Unknown Tokenizer"
+        if filename is None:
+            filename = self.filename
+        with self.output_folder.open(f"{filename}.metadata") as f:
+            if token_count == -1:
+                token_count = self.write_idx
+            f.write("\n".join([tokenizer_name, str(token_count), humanize.metric(token_count, unit="T")]))
+
 
 class DocumentTokenizer(PipelineStep):
     name = "✍️ Writer"
@@ -126,13 +123,12 @@ def __init__(
         save_filename: str = None,  # if defined, the final output filename will be this
         tokenizer_name: str = "gpt2",  # tokenizer to use, from HF
         eos_token: str = "
",  # whether to add the EOS token after each document
-        save_loss_metadata: bool = True,  # save the loss information
+        save_loss_metadata: bool = False,  # save the loss information
         shuffle: bool = True,  # whether to shuffle documents in the dataset,
         seed: int = None,
-        *args,
-        **kwargs,
+        save_final_metadata: bool = True,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__()
         self.output_folder = output_folder
         self.save_filename = save_filename
         self.tokenizer_name = tokenizer_name
@@ -141,6 +137,7 @@ def __init__(
         self.shuffle = shuffle
         self._tokenizer = None
         self.rand = default_rng(seed)
+        self.save_final_metadata = save_final_metadata
 
     def set_up_dl_locks(self, dl_lock, up_lock):
         self.output_folder.set_lock(up_lock)
@@ -160,11 +157,12 @@ def get_loss_values(self, document: Document, encoded):
         return loss_values
 
     def write_unshuffled(self, data, filename):
-        with TokenizedFile(
+        unshuff = TokenizedFile(
             self.output_folder, filename, save_index=not self.shuffle, save_loss_metadata=self.save_loss_metadata
-        ) as unshuff:
-            # tokenize each document's text and write its tokens sequentially to the output .ds
-            for document in data:
+        )
+        # tokenize each document's text and write its tokens sequentially to the output .ds
+        for document in data:
+            with self.track_time():
                 encoded: Encoding = self.tokenizer.encode(document.content)
                 tokens = encoded.ids
                 # loss values
@@ -174,20 +172,26 @@ def write_unshuffled(self, data, filename):
                 # write bytes to disk
                 unshuff.write(tokens, loss_values)
                 # save stats
-                self.stat_update("tokens", len(tokens))
+                self.stat_update("tokens", value=len(tokens))
         return unshuff
 
     def get_output_filename(self, rank, name):
         return "_".join([x for x in [self.save_filename, f"{rank:05d}", f"{name}.ds"] if x])
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         unshuf_filename = self.get_output_filename(rank, "unshuffled")
         logger.info(f'Tokenizing in "{unshuf_filename}"...')
-        unshuffled_file: TokenizedFile = self.write_unshuffled(data, unshuf_filename)
+        outputfile: TokenizedFile = self.write_unshuffled(data, unshuf_filename)
         if self.shuffle:
             shuffled_filename = self.get_output_filename(rank, "shuffled")
-            unshuffled_file.copy(shuffled_filename, self.rand.permutation(len(unshuffled_file.doc_ends)))
-            unshuffled_file.cleanup()
+            # get new TokenizedFile, shuffling docs from original one
+            new_outputfile = outputfile.copy(shuffled_filename, self.rand.permutation(len(outputfile.doc_ends)))
+            # remove and replace original file
+            outputfile.cleanup()
+            outputfile = new_outputfile
+        outputfile.close()
+        if self.save_final_metadata:
+            outputfile.save_final_metadata(self.tokenizer_name)
         self.output_folder.close()
 
     @property
diff --git a/src/datatrove/pipeline/writers/disk_base.py b/src/datatrove/pipeline/writers/disk_base.py
index f5b51dea..4ac8a9cd 100644
--- a/src/datatrove/pipeline/writers/disk_base.py
+++ b/src/datatrove/pipeline/writers/disk_base.py
@@ -2,7 +2,7 @@
 from string import Template
 
 from datatrove.data import Document, DocumentsPipeline
-from datatrove.io import BaseOutputDataFolder, OutputDataFile
+from datatrove.io import BaseOutputDataFile, BaseOutputDataFolder
 from datatrove.pipeline.base import PipelineStep
 from datatrove.utils.typeshelper import StatHints
 
@@ -11,8 +11,8 @@ class DiskWriter(PipelineStep, ABC):
     default_output_filename: str = None
     type = "💽 - WRITER"
 
-    def __init__(self, output_folder: BaseOutputDataFolder, output_filename: str = None, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, output_folder: BaseOutputDataFolder, output_filename: str = None):
+        super().__init__()
         self.output_folder = output_folder
         self.output_filename = Template(output_filename or self.default_output_filename)
 
@@ -25,7 +25,7 @@ def close(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def _get_output_filename(self, document: Document, rank: int = 0, **kwargs):
+    def _get_output_filename(self, document: Document, rank: int | str = 0, **kwargs):
         return self.output_filename.substitute(
             {"rank": str(rank).zfill(5), "data_id": document.data_id, **document.metadata, **kwargs}
         )
@@ -41,16 +41,14 @@ def open(self, output_filename):
         return self.output_folder.open(output_filename)
 
     def write(self, document: Document, rank: int = 0, **kwargs):
-        output_filename = self._get_output_filename(document, rank, **kwargs)
-        output_file: OutputDataFile = self.open(output_filename)
+        output_file: BaseOutputDataFile = self.open(self._get_output_filename(document, rank, **kwargs))
         self._write(document, output_file)
-        output_file.nr_documents += 1
-        self.stat_update(output_filename)
+        self.stat_update(self._get_output_filename(document, "XXXXX", **kwargs))
         self.stat_update(StatHints.total)
 
-    def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
+    def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         with self:
             for document in data:
-                with self.stats.time_manager:
+                with self.track_time():
                     self.write(document, rank)
                 yield document
diff --git a/src/datatrove/pipeline/writers/jsonl.py b/src/datatrove/pipeline/writers/jsonl.py
index 41d9e483..d75002a8 100644
--- a/src/datatrove/pipeline/writers/jsonl.py
+++ b/src/datatrove/pipeline/writers/jsonl.py
@@ -2,7 +2,7 @@
 import json
 
 from datatrove.data import Document
-from datatrove.io import OutputDataFile
+from datatrove.io import BaseOutputDataFile
 from datatrove.pipeline.writers.disk_base import DiskWriter
 
 
@@ -13,8 +13,8 @@ class JsonlWriter(DiskWriter):
     def open(self, output_filename):
         return self.output_folder.open(output_filename, mode="wt", gzip=True)
 
-    def _write(self, document: Document, file: OutputDataFile):
-        file.file_handler.write(
+    def _write(self, document: Document, file: BaseOutputDataFile):
+        file.write(
             json.dumps({key: val for key, val in dataclasses.asdict(document).items() if val}, ensure_ascii=False)
             + "\n"
         )
diff --git a/src/datatrove/tools/check_dataset.py b/src/datatrove/tools/check_dataset.py
index 3ef6801f..fdb54b2d 100644
--- a/src/datatrove/tools/check_dataset.py
+++ b/src/datatrove/tools/check_dataset.py
@@ -7,7 +7,7 @@
 from tokenizers import Tokenizer
 from tqdm import tqdm
 
-from datatrove.io import BaseInputDataFolder, InputDataFile
+from datatrove.io import BaseInputDataFile, BaseInputDataFolder
 
 
 parser = argparse.ArgumentParser()
@@ -21,12 +21,12 @@
 """
 
 
-def load_doc_ends(file: InputDataFile):
+def load_doc_ends(file: BaseInputDataFile):
     with file.open_binary() as f:
         return np.frombuffer(f.read(), dtype=np.uint64)
 
 
-def load_input_mmap(file: InputDataFile):
+def load_input_mmap(file: BaseInputDataFile):
     with file.open_binary() as f:
         return mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
 
@@ -46,7 +46,7 @@ def check_dataset(input_folder: BaseInputDataFolder, tokenizer: str = "gpt2", eo
 
     doc_ends = [load_doc_ends(file) for file in datafiles_index]
     token_inputs = list(map(load_input_mmap, datafiles))
-    loss_inputs = list(map(load_input_mmap, datafiles_loss)) if check_loss else None
+    loss_inputs = list(map(load_input_mmap, datafiles_loss)) if check_loss else [None] * len(token_inputs)
     for filei, (file_doc_ends, file_token_inputs, file_loss_inputs) in enumerate(
         zip(doc_ends, token_inputs, loss_inputs)
     ):
diff --git a/src/datatrove/tools/launch_pickled_pipeline.py b/src/datatrove/tools/launch_pickled_pipeline.py
new file mode 100644
index 00000000..abed5945
--- /dev/null
+++ b/src/datatrove/tools/launch_pickled_pipeline.py
@@ -0,0 +1,23 @@
+import argparse
+
+import dill
+
+from datatrove.executor.base import PipelineExecutor
+from datatrove.io import BaseInputDataFile
+
+
+parser = argparse.ArgumentParser("Loads a pickled pipeline executor and launches it.")
+
+parser.add_argument("path", type=str, help="Path to the pickled file (usually a file called executor.pik)")
+
+
+def main():
+    args = parser.parse_args()
+    executor_path = BaseInputDataFile.from_path(args.path)
+    with executor_path.open_binary() as f:
+        executor: PipelineExecutor = dill.load(f)
+    executor.run()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/datatrove/tools/merge_stats.py b/src/datatrove/tools/merge_stats.py
index c9a847df..c17ae268 100644
--- a/src/datatrove/tools/merge_stats.py
+++ b/src/datatrove/tools/merge_stats.py
@@ -4,6 +4,7 @@
 
 from loguru import logger
 
+from datatrove.io import BaseInputDataFolder, BaseOutputDataFile
 from datatrove.utils.stats import PipelineStats
 
 
@@ -22,18 +23,26 @@
 
 
 def main():
-    args = parser.parse_args()
-    stats_path = args.path
-    output_file = os.path.abspath(args.output)
+    args, file_args = parser.parse_known_args()
+    stats_folder = BaseInputDataFolder.from_path(args.path)
+    # output file
+    extra_args = dict(extra_arg.split("=") for extra_arg in file_args)
+    if (
+        "cleanup" in extra_args
+    ):  # dumb workaround to cast this to boolean (otherwise we'd have to use pydantic or smth)
+        extra_args["cleanup"] = extra_args["cleanup"] == "True"
+    path = extra_args.pop("path", os.path.abspath(args.output))
+    output_file = BaseOutputDataFile.from_path(path, **extra_args)
 
     stats = []
-    for file in os.listdir(stats_path):
-        with open(os.path.join(stats_path, file)) as f:
+    for file in stats_folder.list_files():
+        with file.open() as f:
             stats.append(PipelineStats.from_json(json.load(f)))
-    merged = sum(stats)
-    os.makedirs(os.path.dirname(output_file), exist_ok=True)
-    merged.save_to_disk(output_file)
-    logger.info(f"Processing complete. Results saved to {output_file}.")
+    merged = sum(stats, start=PipelineStats())
+    with output_file.open() as f:
+        merged.save_to_disk(f)
+    logger.info(f"Processing complete. Results saved to {output_file.path}.")
+    logger.info(merged)
 
 
 if __name__ == "__main__":
diff --git a/src/datatrove/utils/logging.py b/src/datatrove/utils/logging.py
new file mode 100644
index 00000000..56b2fcba
--- /dev/null
+++ b/src/datatrove/utils/logging.py
@@ -0,0 +1,34 @@
+import random
+import string
+import sys
+from datetime import datetime
+
+from loguru import logger
+
+from datatrove.io import BaseOutputDataFolder
+
+
+def get_timestamp():
+    return datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+
+def get_random_str(length=5):
+    return "".join(random.choice(string.ascii_lowercase) for _ in range(length))
+
+
+def add_task_logger(logging_dir: BaseOutputDataFolder, rank: int, local_rank: int = 0):
+    logger.remove()
+    logger.add(sys.stderr, level="INFO" if local_rank == 0 else "ERROR")
+    logger.add(logging_dir.open(f"logs/task_{rank:05d}.log"), colorize=True, level="DEBUG")
+    logger.info(f"Launching pipeline for {rank=}")
+
+
+def close_task_logger(logging_dir: BaseOutputDataFolder, rank: int):
+    logger.remove()
+    logging_dir.open(f"logs/task_{rank:05d}.log").close()
+    logger.add(sys.stderr)
+
+
+def log_pipeline(pipeline):
+    steps = "\n".join([pipe.__repr__() if callable(pipe) else "Iterable" for pipe in pipeline])
+    logger.info(f"\n--- 🛠️ PIPELINE 🛠\n{steps}")
diff --git a/src/datatrove/utils/stats.py b/src/datatrove/utils/stats.py
index b6c636d6..b9cce95d 100644
--- a/src/datatrove/utils/stats.py
+++ b/src/datatrove/utils/stats.py
@@ -1,205 +1,293 @@
+import datetime
+import itertools
 import json
 import math
-import os
 import time
-from collections import Counter
+from collections import defaultdict
 from dataclasses import dataclass
+from typing import Callable
 
+import humanize
 
-@dataclass
-class OnlineStats:
-    mean: float
-    variance: float
-    standard_deviation: float
-    total_time: float
-    count: int
-    max: float
-    min: float
+from datatrove.io import BaseOutputDataFile
+
+
+INDENT = " " * 4
+
+
+class OnlineStatsDict(defaultdict):
+    """
+    Stores multiple stats
+    """
+
+    def __init__(self, *_, **kwargs):
+        super().__init__(OnlineStats, **kwargs)
+
+    def __add__(self, other):
+        result = OnlineStatsDict()
+        for key, item in itertools.chain(self.items(), other.items()):
+            result[key] += item
+        return result
+
+    def __repr__(self):
+        return ", ".join(f"{key}: {stats}" for key, stats in self.items())
+
+    def to_dict(self):
+        return {a: b.to_dict() for a, b in self.items()}
 
 
 class Stats:
-    def __init__(self, name: str, indentation: int = 4):
+    def __init__(self, name: str):
         self.name = name
-        self.time_manager = TimeStatsManager()
-        self.counter = Counter()
-        self.doc_len = ComputeOnlineStats()
-        self.indentation = indentation
+        self.time_stats = OnlineTimingStats()
+        self.stats = OnlineStatsDict()
+        self.doc_len_stats = OnlineStats()
 
-    @property
-    def time(self):
-        return self.time_manager.get_stats()
+    def __getitem__(self, item: str) -> "OnlineStats":
+        return self.stats[item]
 
-    @property
-    def length(self):
-        return self.doc_len.get_stats()
+    def __setitem__(self, key: str, value: "OnlineStats"):
+        self.stats[key] = value
 
     def __add__(self, stat):
-        assert self.name == stat.name, f"Summing different blocks {self.name} != {stat.name}"
-        self.counter += stat.counter
-        self.time_manager += stat.time_manager
-        return self
-
-    def _get_frac(self, total_time):
-        return (self.time.total_time / total_time) if total_time > 0 else 0
-
-    def _counter_repr(self):
-        if self.counter == Counter():
-            return ""
-        return f"{' ' * self.indentation + '[' + ''.join([f'{k}={v} ' for k, v in self.counter.items()])}]\n"
-
-    def _time_repr(self, total_time: float):
-        if self.time.total_time == 0:
-            return f"{' ' * self.indentation}Time not computed\n"
-        return (
-            f"{' ' * self.indentation}[{self.time.total_time:.4f}s {self._get_frac(total_time):.2%}"
-            f" {self.time.mean:.4f}±{self.time.standard_deviation:.4f}s/doc]\n"
-        )
+        assert self.name == stat.name, f"Can not merge stats from different blocks {self.name} != {stat.name}"
+        result = Stats(self.name)
+        result.time_stats = self.time_stats + stat.time_stats
+        result.stats = self.stats + stat.stats
+        result.doc_len_stats = self.doc_len_stats + stat.doc_len_stats
+        return result
 
-    def _len_repr(self):
-        if self.doc_len.counter["n"] == 0:
-            return ""
-        return (
-            f"{' ' * self.indentation}[max={self.length.max}, min={self.length.min}"
-            f" {self.length.mean:.0f}±{self.length.standard_deviation:.0f}chars/doc]\n"
+    def __repr__(self, total_time: float = 0.0):
+        return f"\n{INDENT}".join(
+            filter(
+                lambda x: x is not None,
+                [
+                    f"{self.name}",
+                    f"Runtime: {self.time_stats.get_repr(total_time)}" if self.time_stats.total > 0 else None,
+                    f"Stats: {{{self.stats}}}" if len(self.stats) > 0 else None,
+                    f"Doc lens: {self.doc_len_stats}" if self.doc_len_stats.total > 0 else None,
+                ],
+            )
         )
 
-    def __repr__(self, total_time: float = 0.0):
-        return f"{self.name}\n" f"{self._time_repr(total_time)}" f"{self._counter_repr()} " f"{self._len_repr()} "
-
-    def to_dict(self, total_time: float = 0.0):
-        data = {"name": self.name, "stats": self.counter}
-        if self.time.total_time != 0:
-            data["time"] = {
-                "total": self.time.total_time,
-                "n": self.time.count,
-                "%": self._get_frac(total_time) * 100,
-                "mean": self.time.mean,
-                "std_dev": self.time.standard_deviation,
-            }
-        if self.doc_len.counter["n"] != 0:
-            data["doc_len"] = {
-                "n": self.doc_len.counter["n"],
-                "max": self.length.max,
-                "min": self.length.min,
-                "mean": self.length.mean,
-                "std_dev": self.length.standard_deviation,
-            }
-        return data
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "time_stats": self.time_stats.to_dict(),
+            "stats": self.stats.to_dict(),
+            "doc_len_stats": self.doc_len_stats.to_dict(),
+        }
 
     def to_json(self):
         return json.dumps(self.to_dict(), indent=4)
 
+    def save_to_disk(self, file: BaseOutputDataFile | str):
+        dump_json_to_outputfile(self.to_json(), file)
+
     @classmethod
     def from_dict(cls, data):
         stats = cls(data["name"])
-        stats.counter = Counter(data["stats"])
-        if "time" in data:
-            stats.time_manager.running_mean = data["time"]["mean"]
-            stats.time_manager.running_variance = data["time"]["std_dev"]
-            stats.time_manager.counter["total"] = data["time"]["total"]
-            stats.time_manager.counter["n"] = data["time"]["n"]
-        if "doc_len" in data:
-            stats.doc_len.counter["n"] = data["doc_len"]["n"]
-            stats.doc_len.running_mean = data["doc_len"]["mean"]
-            stats.doc_len.running_variance = data["doc_len"]["std_dev"]
-            stats.doc_len.min_value = data["doc_len"]["min"]
-            stats.doc_len.max_value = data["doc_len"]["max"]
+        stats.time_stats = OnlineTimingStats.from_dict(data["time_stats"])
+        stats.stats = OnlineStatsDict()
+        stats.stats.update(data["stats"])
+        stats.doc_len_stats = OnlineStats.from_dict(data["doc_len_stats"])
         return stats
 
 
 class PipelineStats:
-    def __init__(self, stats: list[Stats]):
-        self.stats = stats
+    def __init__(self, stats: list[Stats | Callable] = None):
+        self.stats: list[Stats] = stats if stats else []
+        if self.stats and not isinstance(self.stats[0], Stats):
+            self.stats: list[Stats] = [
+                pipeline_step.stats for pipeline_step in self.stats if hasattr(pipeline_step, "stats")
+            ]
 
     def __add__(self, pipestat):
-        if not isinstance(pipestat, int):
-            self.stats = [x + y for x, y in zip(self.stats, pipestat.stats)]
-        return self
-
-    def __radd__(self, other):
-        return self.__add__(other)
+        if not self.stats:
+            return PipelineStats(pipestat.stats)
+        return PipelineStats([x + y for x, y in zip(self.stats, pipestat.stats)])
 
     @property
     def total_time(self):
-        return sum([stat.time.total_time for stat in self.stats])
+        return sum([stat.time_stats.total for stat in self.stats])
 
-    def __repr__(self):
-        x = f"\n\n{'📉' * 3} STATS {'📉' * 3}\n\n"
-        x += "".join([stat.__repr__(self.total_time) for stat in self.stats])
+    def get_repr(self, text=None):
+        x = f"\n\n{'📉' * 3} Stats{': ' + text if text else ''} {'📉' * 3}\n\n"
+        total_time = self.total_time
+        x += "\n".join([stat.__repr__(total_time) for stat in self.stats])
         return x
 
+    def __repr__(self):
+        return self.get_repr()
+
     def to_json(self):
-        return json.dumps([stat.to_dict(self.total_time) for stat in self.stats], indent=4)
+        return json.dumps([stat.to_dict() for stat in self.stats], indent=4)
 
     @classmethod
     def from_json(cls, data):
         return PipelineStats([Stats.from_dict(stat) for stat in data])
 
-    def save_to_disk(self, path: str):
-        os.makedirs(os.path.dirname(path), exist_ok=True)
-        with open(path, "w") as f:
-            f.write(self.to_json())
+    def save_to_disk(self, file: BaseOutputDataFile | str):
+        dump_json_to_outputfile(self.to_json(), file)
 
 
-class ComputeOnlineStats:
-    def __init__(self):
-        self.counter = Counter()
-        self.running_mean, self.running_variance = 0, 0
-        self.max_value, self.min_value = 0.0, float("inf")
+@dataclass
+class OnlineStats:
+    total: float = 0
+    n: int = 0
+    mean: float = 0.0
+    min: float = float("inf")
+    max: float = float("-inf")
+    _running_variance: float = 0.0
+    unit: str = "doc"
+
+    def update(self, x: float, unit: str = None):
+        if unit:
+            self.unit = unit
+        self.total += x
+        self.n += 1
+
+        self.min = min(self.min, x)
+        self.max = max(self.max, x)
+
+        # https://en.wikipedia.org/wiki/Algorithms_.for_calculating_variance#Welford%27s_online_algorithm
+        delta = x - self.mean
+        self.mean += delta / self.n
+        if self.n > 1:
+            self._running_variance += delta * (x - self.mean)
 
-    def update(self, x: float):
-        self.counter["total"] += x
-        self.counter["n"] += 1
+    @property
+    def variance(self):
+        return self._running_variance / (self.n - 1) if self.n > 1 else 0.0
 
-        self.min_value = min(self.min_value, x)
-        self.max_value = max(self.max_value, x)
-        previous_running_mean = self.running_mean
-        self.running_mean = previous_running_mean + (x - previous_running_mean) / self.counter["n"]
-        if self.counter["n"] == 1:
-            self.running_variance = 0
-        else:
-            self.running_variance = self.running_variance + (x - previous_running_mean) * (x - self.running_mean) / (
-                self.counter["n"] - 1
-            )
+    @property
+    def standard_deviation(self):
+        return math.sqrt(self.variance)
 
     def __add__(self, other):
-        self.counter["total"] += other.counter["total"]
-        n = self.counter["n"] + other.counter["n"]
-        self.running_mean = (
-            (self.counter["n"] * self.running_mean + other.counter["n"] * other.running_mean) / n if n > 0 else 0
-        )
+        if not isinstance(other, type(self)):
+            other = type(self).from_dict(other)
+        # mean and variance: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
+        n = self.n + other.n
+        mean = 0.0
+        _running_variance = 0.0
+        if self.n + other.n > 0:
+            mean = (self.n * self.mean + other.n * other.mean) / n
+            delta = self.mean - other.mean
+            _running_variance = (
+                self._running_variance + other._running_variance + (delta * delta * self.n * other.n) / n
+            )
 
-        delta = self.running_mean - other.running_mean
-        self.running_variance = (
-            self.running_variance + other.running_variance + self.counter["n"] * other.counter["n"] * delta**2 / n
-            if n > 0
-            else 0
+        total = self.total + other.total
+        M = max(self.max, other.max)
+        m = min(self.min, other.min)
+        return type(self)(
+            total=total,
+            n=n,
+            mean=mean,
+            min=m,
+            max=M,
+            _running_variance=_running_variance,
+            unit=self.unit if self.unit != "doc" else other.unit,
         )
 
-        self.max_value = max(self.max_value, other.max_value)
-        self.min_value = min(self.min_value, other.min_value)
-        self.counter["n"] = n
-        return self
-
-    def get_stats(self) -> OnlineStats:
-        return OnlineStats(
-            mean=self.running_mean,
-            variance=self.running_variance,
-            standard_deviation=math.sqrt(self.running_variance),
-            total_time=self.counter["total"],
-            count=self.counter["n"],
-            max=self.max_value,
-            min=self.min_value,
-        )
+    def to_dict(self):
+        if self.total == 0:
+            return 0
+        data = {
+            "total": self.total,
+        }
+        # only relevant if > 1 and we didn't just add 1 all the time
+        if self.n > 1 and self.n != self.total:
+            data["n"] = self.n
+        if self.mean != 1:
+            data["mean"] = self.mean
+        # are there actually different values
+        if self.mean != self.max or self.mean != self.min:
+            data["mean"] = self.mean
+            data["variance"] = self.variance
+            data["std_dev"] = self.standard_deviation
+            data["min"] = self.min
+            data["max"] = self.max
+        if self.unit != "doc":
+            data["unit"] = self.unit
+        return self.total if len(data) == 1 else data
 
+    @classmethod
+    def from_dict(cls, data):
+        if isinstance(data, dict):
+            total = data.get("total")
+            mean = data.get("mean", 1)
+            n = data.get("n", total if mean != 1 else 1)
+            return cls(
+                total=total,
+                n=n,
+                mean=mean,
+                min=data.get("min", mean),
+                max=data.get("max", mean),
+                _running_variance=data.get("variance", 0.0) * (n - 1),
+                unit=data.get("unit", cls.unit),
+            )
+        else:
+            return cls(total=data, min=data, max=data, mean=data, n=data)
+
+    def __repr__(self):
+        if self.mean != 1:
+            # only display relevant fields
+            elements = [
+                (f"min={self.min}, ", self.min != self.mean),
+                (f"max={self.max}, ", self.max != self.mean),
+                (f"{self.mean:.2f}", True),
+                (f"±{self.standard_deviation:.0f}", self.standard_deviation != 0.0),
+            ]
+            return f"{self.total} [" + "".join([t for t, c in elements if c]) + f"/{self.unit}]"
+        return str(self.total)
 
-class TimeStatsManager(ComputeOnlineStats):
-    def __init__(self):
-        super().__init__()
 
+@dataclass
+class OnlineTimingStats(OnlineStats):
     def __enter__(self):
         self._entry_time = time.perf_counter()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        assert self._entry_time is not None
         self.update(time.perf_counter() - self._entry_time)
+
+    def _get_time_frac(self, total_time):
+        return (self.total / total_time) if total_time > 0 else 0
+
+    def get_repr(self, total_time: float = 0.0):
+        if self.total == 0:
+            return "Time not computed"
+        return (
+            f"[{humanize.precisedelta(self.total)} ({self._get_time_frac(total_time):.2%}), "
+            f"{humanize.precisedelta(datetime.timedelta(seconds=self.mean), minimum_unit='milliseconds')}"
+            f"±{humanize.precisedelta(datetime.timedelta(seconds=self.standard_deviation), minimum_unit='milliseconds')}/{self.unit}]"
+        )
+
+    def __repr__(self):
+        return self.get_repr()
+
+    def to_dict(self):
+        data = super().to_dict()
+        if isinstance(data, dict):
+            data["total_human"] = humanize.precisedelta(self.total)
+            data["mean_human"] = humanize.precisedelta(
+                datetime.timedelta(seconds=self.mean), minimum_unit="milliseconds"
+            )
+            data["std_dev_human"] = humanize.precisedelta(
+                datetime.timedelta(seconds=self.standard_deviation), minimum_unit="milliseconds"
+            )
+            data["min_human"] = humanize.precisedelta(
+                datetime.timedelta(seconds=self.min), minimum_unit="milliseconds"
+            )
+            data["max_human"] = humanize.precisedelta(
+                datetime.timedelta(seconds=self.max), minimum_unit="milliseconds"
+            )
+        return data
+
+
+def dump_json_to_outputfile(json_data, file: BaseOutputDataFile | str):
+    if isinstance(file, str):
+        file = BaseOutputDataFile.from_path(file)
+    file.write(json_data)
+    file.close()


Test patch:
diff --git a/tests/pipeline/test_minhash.py b/tests/pipeline/test_minhash.py
index 9000e71b..fc5055ed 100644
--- a/tests/pipeline/test_minhash.py
+++ b/tests/pipeline/test_minhash.py
@@ -126,7 +126,7 @@ def test_buckets_and_cluster(self):
 
             # test duplicate pairs
             for b in range(config.num_buckets):
-                buckets_block(None, bucket=b, world_size=config.num_buckets)
+                buckets_block(None, rank=b, world_size=config.num_buckets)
             bucket_results_folder = LocalInputDataFolder(buckets_folder)
             dup_files = bucket_results_folder.list_files(extension=".dups")
             pairs = defaultdict(set)