Add `Checkpointer` step #1114

plaguss · 2025-01-28T11:03:36Z

Description

Adds a first version of Checkpointer step to write data to the hub while the pipeline executes, useful for longer pipelines.

from datasets import Dataset

from distilabel.pipeline import Pipeline
from distilabel.steps import Checkpointer

dataset = Dataset.from_dict({"a": [1, 2, 3, 4] * 50, "b": [5, 6, 7, 8] * 50})

from distilabel.steps.base import Step, StepInput
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from distilabel.typing import StepOutput

class DoNothing(Step):
    def process(self, *inputs: StepInput) -> "StepOutput":
        for input in inputs:
            yield input

with Pipeline(name="pipeline-with-checkpoints") as pipeline:
    text_generation = DoNothing(
        input_batch_size=60
    )
    checkpoint = Checkpointer(
        repo_id="plaguss/streaming_test_1",
        private=True,
        input_batch_size=50
    )
    text_generation >> checkpoint


if __name__ == "__main__":
    distiset = pipeline.run(
        dataset=dataset,
        use_cache=False
    )
    distiset.push_to_hub(repo_id="plaguss/streaming_test")

logs for the sample pipeline:

[01/28/25 13:37:43] INFO     ['distilabel.pipeline'] 📝 Pipeline data will be written to                                                                                                           base.py:1033
                             '/Users/agus/.cache/distilabel/pipelines/simple-text-generation-pipeline/a2ac64d816f67e4080cfe18d32edb52dc7310f04/executions/dc6b0ffa74e408501f561b8d5c5df2f4d983a5b1             
                             /data/steps_outputs'                                                                                                                                                              
                    INFO     ['distilabel.pipeline'] ⌛ The steps of the pipeline will be loaded in stages:                                                                                        base.py:1065
                              * Legend: 🚰 GeneratorStep 🌐 GlobalStep 🔄 Step                                                                                                                                 
                              * Stage 0:                                                                                                                                                                       
                                - 🚰 'load_data_from_hub_0'                                                                                                                                                    
                                - 🔄 'do_nothing_0'                                                                                                                                                            
                                - 🔄 'checkpointer_0'                                                                                                                                                          
[01/28/25 13:37:44] INFO     ['distilabel.pipeline'] ⏳ Waiting for all the steps of stage 0 to load...                                                                                            base.py:1401
[01/28/25 13:37:45] INFO     ['distilabel.step.checkpointer_0'] Creating repo plaguss/streaming_test_1                                                                                       checkpointer.py:38
[01/28/25 13:37:46] INFO     ['distilabel.pipeline'] ⏳ Steps from stage 0 loaded: 3/3                                                                                                             base.py:1437
                              * 'do_nothing_0' replicas: 1/1                                                                                                                                                   
                              * 'checkpointer_0' replicas: 1/1                                                                                                                                                 
                              * 'load_data_from_hub_0' replicas: 1/1                                                                                                                                           
                    INFO     ['distilabel.pipeline'] ✅ All the steps from stage 0 have been loaded!                                                                                               base.py:1441
                    INFO     ['distilabel.step.load_data_from_hub_0'] 🚰 Starting yielding batches from generator step 'load_data_from_hub_0'. Offset: 0                                    step_wrapper.py:179
                    INFO     ['distilabel.step.load_data_from_hub_0'] 📨 Step 'load_data_from_hub_0' sending batch 0 to output queue                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.load_data_from_hub_0'] 📨 Step 'load_data_from_hub_0' sending batch 1 to output queue                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.do_nothing_0'] 📦 Processing batch 0 in 'do_nothing_0' (replica ID: 0)                                                                       step_wrapper.py:229
                    INFO     ['distilabel.step.do_nothing_0'] 📨 Step 'do_nothing_0' sending batch 0 to output queue                                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.checkpointer_0'] 📦 Processing batch 0 in 'checkpointer_0' (replica ID: 0)                                                                   step_wrapper.py:229
                    INFO     ['distilabel.step.load_data_from_hub_0'] 📨 Step 'load_data_from_hub_0' sending batch 2 to output queue                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.do_nothing_0'] 📦 Processing batch 1 in 'do_nothing_0' (replica ID: 0)                                                                       step_wrapper.py:229
                    INFO     ['distilabel.step.do_nothing_0'] 📨 Step 'do_nothing_0' sending batch 1 to output queue                                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.load_data_from_hub_0'] 📨 Step 'load_data_from_hub_0' sending batch 3 to output queue                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.load_data_from_hub_0'] 🏁 Finished running step 'load_data_from_hub_0' (replica ID: 0)                                                       step_wrapper.py:129
                    INFO     ['distilabel.step.do_nothing_0'] 📦 Processing batch 2 in 'do_nothing_0' (replica ID: 0)                                                                       step_wrapper.py:229
                    INFO     ['distilabel.step.do_nothing_0'] 📨 Step 'do_nothing_0' sending batch 2 to output queue                                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.do_nothing_0'] 📦 Processing batch 3 in 'do_nothing_0' (replica ID: 0)                                                                       step_wrapper.py:229
                    INFO     ['distilabel.step.do_nothing_0'] 📨 Step 'do_nothing_0' sending batch 3 to output queue                                                                        step_wrapper.py:289
                    INFO     ['distilabel.step.do_nothing_0'] 🏁 Finished running step 'do_nothing_0' (replica ID: 0)                                                                       step_wrapper.py:129
[01/28/25 13:37:47] INFO     ['distilabel.step.checkpointer_0'] Uploaded checkpoint 0-0                                                                                                      checkpointer.py:56
                    INFO     ['distilabel.step.checkpointer_0'] 📨 Step 'checkpointer_0' sending batch 0 to output queue                                                                    step_wrapper.py:289
                    INFO     ['distilabel.step.checkpointer_0'] 📦 Processing batch 1 in 'checkpointer_0' (replica ID: 0)                                                                   step_wrapper.py:229
                    INFO     ['distilabel.step.checkpointer_0'] Uploaded checkpoint 0-1                                                                                                      checkpointer.py:56
                    INFO     ['distilabel.step.checkpointer_0'] 📨 Step 'checkpointer_0' sending batch 1 to output queue                                                                    step_wrapper.py:289
                    INFO     ['distilabel.step.checkpointer_0'] 📦 Processing batch 2 in 'checkpointer_0' (replica ID: 0)                                                                   step_wrapper.py:229
[01/28/25 13:37:48] INFO     ['distilabel.step.checkpointer_0'] Uploaded checkpoint 0-2                                                                                                      checkpointer.py:56
                    INFO     ['distilabel.step.checkpointer_0'] 📨 Step 'checkpointer_0' sending batch 2 to output queue                                                                    step_wrapper.py:289
                    INFO     ['distilabel.step.checkpointer_0'] 📦 Processing batch 3 in 'checkpointer_0' (replica ID: 0)                                                                   step_wrapper.py:229
                    INFO     ['distilabel.step.checkpointer_0'] Uploaded checkpoint 0-3                                                                                                      checkpointer.py:56
                    INFO     ['distilabel.step.checkpointer_0'] 📨 Step 'checkpointer_0' sending batch 3 to output queue                                                                    step_wrapper.py:289
                    INFO     ['distilabel.step.checkpointer_0'] 🏁 Finished running step 'checkpointer_0' (replica ID: 0)                                                                   step_wrapper.py:129
Generating train split: 200 examples [00:00, 107353.57 examples/s]
Creating parquet from Arrow format: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 1449.31ba/s]
Uploading the dataset shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.31it/s]

… specified points during the life of the pipeline

…frequency of file writes

github-actions · 2025-01-28T11:05:13Z

Documentation for this PR has been built. You can view it at: https://distilabel.argilla.io/pr-1114/

codspeed-hq · 2025-01-28T11:07:57Z

CodSpeed Performance Report

Merging #1114 will degrade performances by 35.85%

_{Comparing stream-push-to-hub (61f3387) with develop (84ea198)}

Summary

❌ 1 regressions

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	`BASE`	`HEAD`	Change
❌	`test_cache_time`	557.5 ms	869.2 ms	-35.85%

… the HfAPI

…ements in the inputs

gabrielmbmb

LGTM! just some comments

gabrielmbmb · 2025-01-28T17:10:58Z

src/distilabel/steps/checkpointer.py

+from huggingface_hub import HfApi
+
+
+class Checkpointer(Step):


Maybe it's better if we call it HuggingFaceHubCheckpointer?

gabrielmbmb · 2025-01-28T17:11:30Z

src/distilabel/steps/checkpointer.py

+        if not self._api.repo_exists(repo_id=self.repo_id, repo_type="dataset"):
+            self._logger.info(f"Creating repo {self.repo_id}")
+            self._api.create_repo(
+                repo_id=self.repo_id, repo_type="dataset", private=self.private
+            )


let's move this logic to a method that can be unit tested

gabrielmbmb · 2025-01-28T17:22:12Z

src/distilabel/steps/checkpointer.py

+        for i, input in enumerate(inputs):
+            # Each section of *inputs corresponds to a different configuration of the pipeline
+            with tempfile.NamedTemporaryFile(
+                mode="w", suffix=".jsonl", delete=False


why delete=False instead of letting the context manager handle that?

plaguss added 5 commits January 28, 2025 11:58

Add special step Checkpointer to read data from the cache and push at…

e161204

… specified points during the life of the pipeline

Redirect import

75f12d3

Add extra guard for the checkpoint step

5e1a4e8

Set data dir in the checkpointer step and allow configuration on the …

9af338c

…frequency of file writes

Add argument to allow modifying the dump batch size from the pipeline

29959f8

plaguss requested a review from gabrielmbmb January 28, 2025 11:03

plaguss added 4 commits January 28, 2025 13:40

Update Checkpointer to just push the current batch to a dataset using…

7528815

… the HfAPI

Push different config files depending on the number of leaf steps (el…

3b64892

…ements in the inputs

Add docs for the Checkpointer step

34cb984

Add emoji to locate the checkpoint

8f1683e

gabrielmbmb requested changes Jan 28, 2025

View reviewed changes

plaguss added 2 commits January 28, 2025 19:24

Rename checkpointer and let the context manager remove the file

58b6f89

Merge branch 'develop' into stream-push-to-hub

61f3387

plaguss marked this pull request as ready for review January 28, 2025 19:05

plaguss merged commit f5ddbc6 into develop Jan 28, 2025
6 of 7 checks passed

plaguss deleted the stream-push-to-hub branch January 28, 2025 19:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `Checkpointer` step #1114

Add `Checkpointer` step #1114

plaguss commented Jan 28, 2025 •

edited

Loading

github-actions bot commented Jan 28, 2025

codspeed-hq bot commented Jan 28, 2025 •

edited

Loading

gabrielmbmb left a comment

gabrielmbmb Jan 28, 2025

gabrielmbmb Jan 28, 2025

gabrielmbmb Jan 28, 2025

Add Checkpointer step #1114

Add Checkpointer step #1114

Conversation

plaguss commented Jan 28, 2025 • edited Loading

Description

github-actions bot commented Jan 28, 2025

codspeed-hq bot commented Jan 28, 2025 • edited Loading

CodSpeed Performance Report

Merging #1114 will degrade performances by 35.85%

Summary

Benchmarks breakdown

gabrielmbmb left a comment

Choose a reason for hiding this comment

gabrielmbmb Jan 28, 2025

Choose a reason for hiding this comment

gabrielmbmb Jan 28, 2025

Choose a reason for hiding this comment

gabrielmbmb Jan 28, 2025

Choose a reason for hiding this comment

Add `Checkpointer` step #1114

Add `Checkpointer` step #1114

plaguss commented Jan 28, 2025 •

edited

Loading

codspeed-hq bot commented Jan 28, 2025 •

edited

Loading