Partitioned Processing with Checkpointing

This document describes DataJuicer’s fault-tolerant processing system with partitioning, checkpointing, and event logging.

Overview

The ray_partitioned executor splits datasets into partitions and processes them with configurable checkpointing. Failed jobs can resume from the last checkpoint.

Checkpointing strategies:

• every_n_ops - checkpoint every N operations (default, balanced)

• every_op - checkpoint after every operation (max protection, impacts performance)

• manual - checkpoint only after specified operations (best for known expensive ops)

• disabled - no checkpointing (best performance)

Directory Structure

{work_dir}/{job_id}/
├── job_summary.json              # Job metadata (created on completion)
├── events_{timestamp}.jsonl      # Machine-readable event log
├── dag_execution_plan.json       # DAG execution plan
├── checkpoints/                  # Checkpoint data
│   ├── partitioning_info.json    # Saved row boundaries and partition hashes
│   └── checkpoint_op_*.parquet/  # Per-operation partition checkpoints
├── partitions/                   # Input partitions
├── logs/                         # Human-readable logs
└── metadata/                     # Job metadata

Configuration

Partition Modes

Auto mode (recommended) - analyzes data and resources to determine optimal partitioning:

executor_type: ray_partitioned

partition:
  mode: "auto"
  max_concurrent_partitions: "auto"  # Resource-aware driver concurrency
  target_size_mb: 256    # Target partition size (128, 256, 512, or 1024)
  size: 5000             # Fallback if auto-analysis fails
  max_size_mb: 256       # Fallback max size

Manual mode - specify exact partition count:

partition:
  mode: "manual"
  num_of_partitions: 8
  max_concurrent_partitions: "auto"

max_concurrent_partitions: "auto" is the default. It is resolved after operator resource planning: GPU pipelines use the tightest CPU/GPU worker capacity reported by the Ray cluster, while CPU-only pipelines use a conservative outer-pipeline cap of 4. The actual concurrency is also bounded by the partition count and any explicit global actor num_proc budget. Set a positive integer to override the automatic limit.

Checkpointing

checkpoint:
  enabled: true
  strategy: every_n_ops  # every_n_ops (default), every_op, manual, disabled
  n_ops: 5               # Default: checkpoint every 5 operations
  op_names:              # For manual strategy - checkpoint after expensive ops
    - document_deduplicator
    - embedding_mapper

When checkpointing is enabled, the initial run saves checkpoints/partitioning_info.json. For every logical partition, this file records:

• start_row (inclusive) and end_row (exclusive) in the ordered input;

• the partition row count;

• a stable hash of the complete partition contents;

• the partition hash algorithm used by the writer.

These values allow an explicit resume to recreate the original logical partitions even when Ray produces a different physical block layout in the new process. Complete partition hashes are sensitive to row order, independent of Ray batch boundaries, and validated before any checkpoint is reused.

Intermediate Storage

intermediate_storage:
  format: "parquet"              # parquet, arrow, jsonl
  compression: "snappy"          # snappy, gzip, none
  preserve_intermediate_data: true
  retention_policy: "keep_all"   # keep_all, keep_failed_only, cleanup_all

Usage

Running Jobs

# Auto partition mode
dj-process --config config.yaml --partition.mode auto

# Manual partition mode
dj-process --config config.yaml --partition.mode manual --partition.num_of_partitions 4

# Optional: start a new job with a custom job ID
dj-process --config config.yaml --job_id my_experiment_001

When checkpointing is enabled, a new ray_partitioned job prints a resume token:

Resume token: 20260805_115141_81270d. Rerun the original command with
--resume 20260805_115141_81270d to resume this job.

Resuming Jobs

# Use the token printed by the initial run. Keep the input and recipe unchanged.
dj-process --config config.yaml --resume 20260805_115141_81270d

# A custom ID from the initial run can be used in the same way.
dj-process --config config.yaml --resume my_experiment_001

--resume is supported only by the ray_partitioned executor. It performs a strict resume in the following order:

1. locate the original work and checkpoint directories;

2. verify that the current configuration matches the original run;

3. load the saved partition count, row boundaries, and content hashes;

4. recreate the original partitions with the saved row boundaries;

5. validate every complete partition hash;

6. load completed checkpoints and process only the unfinished work.

If metadata is missing, row boundaries are invalid, the input has changed, or a content hash does not match, explicit resume stops with an error and leaves the existing checkpoints unchanged.

--job_id remains available for custom job naming and backward compatibility. For fault-tolerant continuation, prefer --resume: the legacy --job_id resumption path keeps its previous behavior and may clear mismatched checkpoints before starting fresh. Metadata created by an older Data-Juicer version can still be read, but it does not contain the row boundaries and full content hashes required by explicit resume; --resume therefore rejects it without deleting its checkpoints.

If both arguments are supplied, their values must be identical:

dj-process --config config.yaml \
  --job_id my_experiment_001 \
  --resume my_experiment_001

Checkpoint Strategies

# Every operation
dj-process --config config.yaml --checkpoint.strategy every_op

# Every N operations
dj-process --config config.yaml --checkpoint.strategy every_n_ops --checkpoint.n_ops 3

# Manual
dj-process --config config.yaml --checkpoint.strategy manual --checkpoint.op_names op1,op2

Auto-Configuration

In auto mode, the optimizer:

1. Samples the dataset to detect modality (text, image, audio, video, multimodal)

2. Measures memory usage per sample

3. Analyzes pipeline complexity

4. Calculates partition size targeting the configured target_size_mb

Default partition sizes by modality:

Modality	Default Size	Max Size	Memory Multiplier
Text	10000	50000	1.0x
Image	2000	10000	5.0x
Audio	1000	4000	8.0x
Video	400	2000	20.0x
Multimodal	1600	6000	10.0x

Job Management Utilities

Monitor

# Show progress
python -m data_juicer.utils.job.monitor {job_id}

# Detailed view
python -m data_juicer.utils.job.monitor {job_id} --detailed

# Watch mode
python -m data_juicer.utils.job.monitor {job_id} --watch --interval 10

from data_juicer.utils.job.monitor import show_job_progress

data = show_job_progress("job_id", detailed=True)

Stopper

# Graceful stop
python -m data_juicer.utils.job.stopper {job_id}

# Force stop
python -m data_juicer.utils.job.stopper {job_id} --force

# List running jobs
python -m data_juicer.utils.job.stopper --list

from data_juicer.utils.job.stopper import stop_job

stop_job("job_id", force=True, timeout=60)

Common Utilities

from data_juicer.utils.job.common import JobUtils, list_running_jobs

running_jobs = list_running_jobs()

job_utils = JobUtils("job_id")
summary = job_utils.load_job_summary()
events = job_utils.load_event_logs()

Event Types

• job_start, job_complete, job_failed

• partition_start, partition_complete, partition_failed

• op_start, op_complete, op_failed

• checkpoint_save, checkpoint_load

Performance Considerations

Checkpoint vs Ray Optimization Trade-off

Key insight: Checkpointing interferes with Ray’s automatic optimization.

Ray optimizes execution by fusing operations together and pipelining data. Each checkpoint forces materialization, which breaks the optimization window:

Without checkpoints:     op1 → op2 → op3 → op4 → op5
                         |___________________________|
                              Ray optimizes entire window

With every_op:           op1 | op2 | op3 | op4 | op5
                         materialize at each | (5 barriers)

With every_n_ops(5):     op1 → op2 → op3 → op4 → op5 |
                         |_____________________________|
                              Ray optimizes all 5 ops

Checkpoint Cost Analysis

Cost Type	Typical Value
Checkpoint write	~2-5 seconds
Cheap op execution	~1-2 seconds
Expensive op execution	minutes to hours

For cheap operations, checkpointing costs MORE than re-running on failure.

Example pipeline analysis:

filter(1s) → mapper(2s) → deduplicator(300s) → filter(1s)

Strategy         | Overhead  | Protection Value
-----------------|-----------|------------------
every_op         | ~20s      | Save 1-304s on failure
after dedup only | ~5s       | Save 300s on failure
disabled         | 0s        | Re-run everything

Strategy Recommendations

Job Duration	Recommended Strategy	Rationale
< 10 min	disabled	Re-running is cheap
10-60 min	every_n_ops (n=5)	Balanced protection
> 60 min with expensive ops	manual	Checkpoint after expensive ops only
Unstable infrastructure	every_n_ops (n=2-3)	Accept overhead for reliability

Operation Categories

Expensive operations (checkpoint after these):

• *_deduplicator - Global state, expensive computation

• *_embedding_* - Model inference

• *_model_* - Model inference

• *_vision_* - Image/video processing

• *_audio_* - Audio processing

Cheap operations (skip checkpointing):

• *_filter - Simple filtering

• clean_* - Text cleaning

• remove_* - Field removal

Storage Recommendations

• Event logs: fast storage (SSD)

• Checkpoints: large capacity storage

• Partitions: local storage

Partition Sizing Trade-offs

• Smaller partitions: better fault tolerance, more scheduling overhead

• Larger partitions: less overhead, coarser recovery granularity

Troubleshooting

Job resumption fails:

ls -la ./outputs/{work_dir}/{job_id}/job_summary.json
ls -la ./outputs/{work_dir}/{job_id}/checkpoints/
cat ./outputs/{work_dir}/{job_id}/checkpoints/partitioning_info.json

Use the resume token printed by the original run and rerun the same recipe with --resume. Check the error log for configuration mismatch, missing partition metadata, invalid row boundaries, or partition content hash mismatch. Explicit resume does not delete checkpoints when validation fails.

Check Ray status:

ray status

View logs:

cat ./outputs/{work_dir}/{job_id}/events_*.jsonl
tail -f ./outputs/{work_dir}/{job_id}/logs/*.txt