Add elastic data sharding with node-local Ray execution

[Qirui-jiao]

Motivation

Processing a large dataset on multiple independent nodes currently requires
users to manually split the input, assign files to nodes, track failures, and
merge the outputs.

This PR adds an elastic sharding example that:

• pre-splits large JSONL datasets into deterministic shards;
• lets multiple nodes dynamically claim shards through shared storage;
• processes every claimed shard with a node-local Data-Juicer Ray executor;
• supports retries, stale-claim recovery, status inspection, and ordered merge;
• provides shared-storage coordination for Worker-broadcast DLC jobs and
  documents the separate Launcher + mpirun topology required by MPIJob.

This differs from a shared multi-node Ray cluster: the shared filesystem
coordinates work between nodes, while Ray is used independently inside each
node.

What does this PR change?

Elastic shard job

Add demos/elastic_sharding/shard_job.py with the following commands:

• prepare

  • validates a Data-Juicer recipe;
  • scans local JSONL input without loading the complete dataset into memory;
  • creates deterministic, non-empty shards;
  • records input fingerprints, recipe hash, execution settings, and shard
    metadata in a manifest.

• worker

  • atomically claims one shard at a time through a shared POSIX filesystem;
  • reclaims stale claims after the configured timeout;
  • runs every shard with executor_type=ray;
  • uses ray_address=local by default, so each node has an independent Ray
    runtime;
  • stores isolated attempt logs and outputs;
  • validates the processed JSONL before publishing completion metadata;
  • verifies the lock inode and claim token under the advisory lock before
    publishing success, so an expired attempt cannot overwrite or beat a
    replacement attempt that already reclaimed the shard.

• status

  • reports pending, running, stale, completed, and failed shards.

• retry

  • archives failed state and requeues selected shards.

• merge

  • verifies row counts and SHA256 checksums;
  • merges successful outputs in the original shard order;
  • atomically publishes the final JSONL.

PAI-DLC entry point

Add a Worker-broadcast dlc mode to
demos/elastic_sharding/dlc_job.py. Keep two_node_test.py as a
backward-compatible strict two-node wrapper.

For DLC job types that execute the configured command on every Worker, the job
only needs one configured startup command:

cd /mnt/data/data-juicer && \
python demos/elastic_sharding/dlc_job.py dlc \
  --job-dir /mnt/shared/data-juicer-jobs/multi-node-job-001 \
  --nodes 4 \
  --num-shards 16 \
  --ray-address local

All Worker instances in that topology run the same entry point. They use the shared job
directory to:

1. elect one preparation coordinator;
2. split the input once;
3. claim and process shards independently;
4. wait for all shards to complete;
5. elect one finalization coordinator;
6. report participating Workers and merge the result.

The Worker-broadcast entry point does not depend on platform-specific rank
environment variables. PAI-DLC MPIJob has a different topology: its configured
command runs only on the Launcher, so the Launcher must use DLC's generated
hostfile and mpirun -np N -npernode 1 to start one shard worker on each GPU
Worker. The MPIJob Workers can still start independent node-local Ray runtimes.

By default, Workers have no per-node claim limit, so the live Workers can
finish all remaining shards even when fewer Workers start than expected.
An optional --require-all-nodes --nodes N strict mode caps claims and verifies
that all N Worker hostnames participated.

The DLC wrapper scopes prepare-result.json, abort.json, phase locks, and
finalize-result.json to one submission generation. It resolves the generation
from an explicit --run-id or from PAI_JOB_ID, DLC_JOB_ID, or JOB_ID,
then stores coordination state below a hash of that identifier. Workers from
the same submission still share idempotent phase results, while a later
submission that reuses the same unchanged job directory no longer inherits a
transient prepare failure, stale abort, or failed finalize result.

Recipe and documentation

• Add a small CPU-only shard-safe Data-Juicer recipe using existing Mapper and
  Filter operators.
• Add a GPU smoke-test recipe and four-row text-pair dataset. One claimed shard
  runs through a CPU whitespace Mapper, GPU sentiment and topic classification
  Mappers, a GPU CLIP similarity Filter, and a CPU length Filter in the
  node-local Ray executor. Every GPU stage uses Ray task mode with
  num_proc: 1 and num_gpus: 1, allowing one GPU per node to be reused
  sequentially.
• Add a single-node four-GPU variant using four Ray input blocks and four
  one-GPU tasks per GPU operator. This exercises all four cards inside one
  claimed shard instead of processing four shards sequentially.
• Add English and Chinese documentation covering:
  • DLC configuration;
  • shared NAS/CPFS requirements;
  • node-local Ray execution;
  • CPU + GPU recipe execution and model-cache requirements;
  • status and output paths;
  • retry and failure handling;
  • current limitations.
• Add the demo to the English and Chinese demo indexes.

Compatibility

This is an additive demo and does not change existing Data-Juicer execution
paths or operator behavior.

The current implementation only accepts shard-independent Mapper and Filter
operators. Dataset-level operations such as deduplication, selection, grouping,
and aggregation are rejected during preparation.

Requirements and limitations

• The job directory must be on a shared POSIX filesystem supporting atomic file
  creation, atomic rename, and fcntl advisory locks, such as NAS/NFS/CPFS.
• All Workers must see the input, media files, code, and job directory at
  consistent paths.
• Inputs must be local JSONL files or directories containing JSONL files.
• Claim expiration currently uses a static timeout without heartbeats.
• ray_address=local starts an independent Ray runtime for each shard attempt.
• Optional strict DLC mode caps work per Worker so all configured Workers must
  participate; default elastic mode has no cap.
• Every Worker-broadcast dlc invocation needs one stable submission identity.
  PAI-DLC normally supplies it through a Job ID environment variable; other
  launchers must pass the same explicit --run-id to every Worker and use a
  new value for each new submission.
• The bundled GPU smoke test requires at least one visible NVIDIA GPU per
  Worker and a CUDA-enabled PyTorch environment. Its three Hugging Face models
  must be pre-downloaded or downloadable by every Worker.
• Data-Juicer/Datasets infers nested JSON schemas independently for each shard.
  When records omit optional nested keys, the merged multi-shard output can
  represent those keys differently from a one-shard run (null versus absent).
  Record membership is preserved, but Ray output order and byte representation
  are not guaranteed across runs. Callers that require byte-identical output
  should normalize the input schema and sort by a stable record key after
  processing.

Testing

Run:

python -m pytest -q tests/demos/test_elastic_sharding.py

Result:

22 passed

The tests cover:

• deterministic splitting and metadata normalization;
• recipe validation;
• exclusive claims and stale-claim recovery;
• rejection of a late success from an expired attempt after a replacement
  attempt has claimed the shard;
• Ray command construction and output materialization;
• retry limits and ordered merge;
• distinct multi-node ownership verification;
• Worker-broadcast DLC coordination and explicit MPIJob topology documentation;
• per-submission isolation of transient prepare failures, stale abort state,
  and transient finalize failures;
• explicit and PAI-DLC-provided submission IDs;
• elastic completion when fewer Workers are live than configured;
• strict participation verification with three Workers;
• validation that the bundled GPU recipe contains a
  CPU -> GPU Mapper -> GPU Mapper -> GPU Filter -> CPU chain and explicitly
  reserves one GPU for every CUDA operator;
• validation that the single-node four-GPU recipe configures four Ray blocks,
  four-way task concurrency, and one GPU per task;
• preparation and Worker failure propagation.

The included Mapper/Filter recipe was also exercised end to end on a CPU-only
host with 64 CPU cores, no visible CUDA devices, Ray 2.55.1, and Data-Juicer
1.5.3:

shards  completed  rows  failed  retries  elapsed
1       1/1        6     0       0        18.768 s
2       2/2        6     0       0        35.505 s
4       4/4        6     0       0        67.924 s

Every completion record reported executor_type=ray and
ray_address=local, and every shard succeeded on its first attempt. The 1-, 2-,
and 4-shard runs preserved all six records in source order and produced
semantically equivalent non-null JSON data.

These timings are correctness smoke-test results, not an acceleration
benchmark. For this tiny input, repeated local Ray startup dominates runtime,
so production shards should be large enough to amortize per-attempt startup
cost.

The separate PAI-DLC PyTorch benchmark harness was also tested with:

python -m pytest -q two_node_test/tests/test_benchmark.py

16 passed, 1 skipped

The skipped test is an opt-in local Ray CPU integration test guarded by
RUN_TWO_NODE_CPU_SMOKE=1. The remaining tests cover deterministic benchmark
input generation, PyTorchJob rank discovery, shared coordination paths,
failure propagation, GPU recipe validation, output validation, telemetry
summaries, report generation, and atomic test coordination.

PAI-DLC GPU scaling validation

The shard workflow was exercised on real PAI-DLC PyTorch jobs without
mpirun. DLC ran the same startup command on each node and supplied RANK,
WORLD_SIZE, MASTER_ADDR, and MASTER_PORT. Each node had 8 NVIDIA H20 GPUs
and ran an independent node-local Ray runtime.

The benchmark recipe used:

• one CPU whitespace normalization Mapper;
• one GPU sentiment Mapper;
• one GPU topic classification Mapper;
• one GPU CLIP text-pair similarity Filter;
• one CPU text-length Filter;
• three pre-downloaded local Hugging Face models;
• num_proc: 8, num_gpus: 1, batch_size: 64, task execution mode, and 64
  Ray input blocks for every GPU stage.

The following table uses the interval from the earliest shard claim to the
latest shard completion for processing time. End-to-end time covers prepare
start through merge completion.

Run	Input rows	1-node processing	2-node processing	Processing speedup	Scaling efficiency	1-node end-to-end	2-node end-to-end	End-to-end speedup
exp03	100,000	438.229 s	257.572 s	1.701x	85.07%	452.021 s	273.509 s	1.653x
exp04	400,000	1,100.678 s	667.217 s	1.650x	82.48%	1,131.243 s	699.532 s	1.617x

Both post-fix comparisons exceeded the benchmark's 1.6x initial scaling
threshold.
For the 400,000-row run, processing throughput increased from 363.412 rows/s
on 8 GPUs to 599.506 rows/s on 16 GPUs, and end-to-end latency fell by 38.16%.

Larger shards also amortized fixed model and Ray task overhead:

Input	1-node average GPU utilization	2-node average GPU utilization	1-node zero-utilization samples	2-node zero-utilization samples
100,000 rows (exp03)	14.71%	13.14%	50.96%	57.53%
400,000 rows (exp04)	23.25%	19.86%	30.13%	37.91%

All 8 GPUs in the single-node runs and all 16 GPUs in the two-node runs
reported nonzero activity. Peak memory remained 1,295 MiB per H20, leaving
substantial room for larger batches. In exp04, increasing the input by 4x
increased processing time by only 2.512x on one node and 2.590x on two nodes;
processing throughput therefore improved by 59.26% and 54.42%, respectively.

Correctness checks passed for every run:

• both modes used identical input and recipe SHA256 fingerprints;
• every shard succeeded on its first attempt and the two-node shards had
  distinct Worker hostnames;
• 100,000-row and 400,000-row outputs had no missing, duplicate, or unexpected
  sample_id values;
• after aligning by sample_id, the complete JSON records from one-node and
  two-node runs were identical;
• differing merged-file SHA256 values were caused only by nondeterministic Ray
  output order.

These are initial scaling results rather than final statistical benchmarks.
Each configuration was observed once per run. The exp04 one-node and two-node
jobs overlapped in time and may have contended for shared storage or underlying
infrastructure. The exp03 and exp04 code commits also differed, and manifests
contained dirty or unavailable Git state. A final benchmark should use one
clean commit, run the one-node and two-node jobs sequentially, and repeat each
configuration at least three times.

The two-node traces also expose startup skew: exp03's second node claimed work
8.024 seconds after the first, while exp04's delay was 19.408 seconds. For
exp04, removing that skew would improve the estimated processing speedup from
1.650x to approximately 1.699x. A pre-claim all-ranks-ready barrier is a useful
follow-up optimization.

Manual validation status

• Unit and coordination tests
• Formatting and static checks
• Real JSONL preparation smoke test
• Single-node, CPU-only, local-Ray end-to-end runs with 1, 2, and 4 shards
• Row-count, record-order, completion-metadata, and semantic-output checks
• One-GPU and four-GPU recipe validation and JSONL preparation
• Real CUDA execution on 8 NVIDIA H20 GPUs
• Multi-node PAI-DLC execution on 2 nodes / 16 NVIDIA H20 GPUs
• 100,000-row post-fix scaling validation
• 400,000-row scaling and full output-equivalence validation

The real DLC runs validate shared-storage visibility, distinct concurrent
Worker ownership, node-local 8-GPU Ray execution, first-attempt completion,
cross-node failure-aware coordination, final merge, and full semantic output
equivalence. Remaining work is performance tuning and a statistically
controlled repeated benchmark, not functional multi-node validation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[fengrui-z]

LGTM