Skip to content

feat: auto-shard eligible multinode jobs - #1041

Open
fengrui-z wants to merge 1 commit into
datajuicer:mainfrom
fengrui-z:agent/core-auto-multinode-sharding
Open

feat: auto-shard eligible multinode jobs#1041
fengrui-z wants to merge 1 commit into
datajuicer:mainfrom
fengrui-z:agent/core-auto-multinode-sharding

Conversation

@fengrui-z

Copy link
Copy Markdown
Collaborator

Summary

This PR promotes the shared-POSIX elastic sharding mechanism from
demos/elastic_sharding into Data-Juicer core and enables it in conservative
auto mode by default.

When dj-process is broadcast to multiple distributed ranks, Data-Juicer can
automatically detect a physical multi-node launch and shard eligible JSONL
pipelines across nodes without requiring a cross-node Ray cluster.

Motivation

Worker-broadcast jobs currently start the same Data-Juicer pipeline on every
process, which may duplicate work and cause multiple processes to write the
same output.

PR #1015 demonstrated that deterministic pre-sharding, shared-filesystem shard
claims, node-local Ray execution, retries, and ordered merge can coordinate
these jobs effectively. This PR moves that mechanism into core and adds
automatic activation and semantic safety checks.

Behavior

The feature is controlled by:

elastic_sharding:
  mode: auto  # auto | on | off

In auto mode, Data-Juicer:

1. Detects Torch, OpenMPI, or Slurm distributed rank metadata.
2. Requires a stable submission ID from configuration or launcher environment.
3. Rendezvous all advertised ranks through shared POSIX storage.
4. Counts distinct hostnames instead of treating WORLD_SIZE as the node count.
5. Elects one shard worker per physical node.
6. Creates deterministic JSONL shards based on node count and target shard size.
7. Dynamically assigns shards through shared-filesystem claims.
8. Runs every shard with an isolated node-local Ray executor.
9. Validates shard row counts and checksums.
10. Publishes one ordered merged JSONL result from rank 0.

Child shard processes explicitly disable elastic sharding to prevent recursive
coordination.

## Conservative eligibility checks

Automatic sharding is enabled only when the pipeline can be proven
partition-safe.

The current implementation requires:

- local JSONL input and output;
- default or node-local ray execution;
- only Mapper and Filter operators;
- no deduplicators or other global operations;
- no Selectors, Groupers, Aggregators, or Pipelines;
- no sampling or parallel/sharded export modes;
- no fixed shared stats_export_path or save_dir;
- custom operators must declare partition_safe = True.

When a confirmed multi-node launch contains an ineligible pipeline, auto
mode runs the original executor only on rank 0 and makes the remaining ranks
wait for the shared result. This prevents duplicate output while preserving
whole-dataset semantics.

mode: on turns unmet prerequisites into errors. mode: off preserves the
legacy per-process behavior.

## Reliability

The core shard state machine provides:

- deterministic input splitting and ordered merge;
- atomic shared-filesystem metadata publication;
- exclusive shard claims;
- heartbeat-renewed leases;
- token fencing so expired workers cannot publish stale results;
- configurable retries;
- per-attempt logs and metadata;
- input, recipe, output, row-count, and SHA256 validation;
- idempotent preparation and final-result recovery.

The coordination filesystem is verified using a POSIX hard-link probe.

## Compatibility

The existing demos/elastic_sharding/shard_job.py entry point remains
available as a thin compatibility wrapper around the core implementation.
Scheduler-specific DLC helpers and the manual
prepare/worker/status/retry/merge workflow continue to work.

## Configuration

The following options are added under elastic_sharding:

- mode
- run_id
- coordination_dir
- rendezvous_timeout_secs
- rendezvous_poll_interval_secs
- result_timeout_secs
- shards_per_node
- target_shard_size_mb
- max_shards
- lock_timeout_secs
- heartbeat_interval_secs
- max_retries
- poll_interval_secs

## Tests

- Core launch-context, stable job ID, factory, rendezvous, hostname
  de-duplication, safety analysis, heartbeat, and coordinator-only tests.

- Existing elastic-sharding demo regression suite.
- Configuration and executor factory suites.
- Real tools/process_data.py compatibility smoke test with
  elastic_sharding.mode=off.

Validation results:

- 158 passed, 1 skipped for config, factory, core sharding, and demo suites.
- 29 passed in the final focused regression run.
- Black, isort, flake8, Python compilation, YAML parsing, and
  git diff --check passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant