Skip to content

Ray execution

Ray is an optional placement backend for the existing NexuML Lightning session. Use it when a scenario should run on an existing Ray cluster. NexuML does not maintain a parallel Ray-specific model or training loop.

Install

uv pip install "nexuml[ray]"

Add nexuml[s3] when Ray run storage or datasets use S3-compatible storage. Add the DALI integration separately when the selected data loader requires it.

The driver and workers need compatible Python/Ray environments. Keep exact environment/image/CUDA choices in the consuming project or cluster configuration, not in ScenarioSpec.

Configure an existing cluster

from nexuml.core.types import RayClusterTarget, RayExecutionSpec

execution = RayExecutionSpec(
    target=RayClusterTarget(
        address="ray://ray.example.org:10001",
        working_dir=".",
        py_executable="uv run --locked python",
    ),
    workers=4,
    resources_per_worker={"CPU": 4, "GPU": 1},
    storage_path="s3://my-bucket/nexuml/runs",
)

Attach it to the scenario and use the normal command:

nexuml train my-scenario

Each worker restores the typed scenario, creates the normal NexuSession, and executes its lifecycle with a Ray-prepared Lightning trainer.

Lightning strategies

training.strategy remains the single source of truth. Ray currently maps auto/ddp, fsdp, and deepspeed to Ray's official Lightning strategy integrations.

TrainingSpec(strategy="ddp")

For strategy-specific settings, reference Ray's real strategy class:

from ray.train.lightning import RayFSDPStrategy

from nexuml import strategy

TrainingSpec(strategy=strategy(RayFSDPStrategy, state_dict_type="full"))

The class and its constructor parameters remain navigable and statically checkable. Required third-party strategy packages still have to exist in the worker environment.

Ray Jobs

NexuML intentionally does not wrap the Ray Jobs lifecycle. If the driver itself should run remotely/detached, use Ray's native CLI:

ray job submit --working-dir . -- \
  uv run --locked nexuml train -c configs/my-scenario.yaml

Use Ray's own status/log/stop tooling for the job lifecycle.

Distributed semantic restrictions

Ray currently rejects two kinds of stateful post-training work rather than silently changing their meaning:

  • scenarios with evaluation.algorithms;
  • components whose definition declares requires_post_train_fit.

Both need global state across rank-sharded data. Until NexuML implements correct cross-worker finalization, keep those scenarios local. Ordinary Lightning validation/test metrics can still be reduced by the distributed trainer.

Shared datasets and S3 WebDataset

A supported shared-data workflow is:

NexuML dataset
  → WebDataset export
  → S3-compatible storage
  → ExportedDataset
  → direct tar streaming and worker-local index staging
  → DALI WebDataset reader

Export to S3 with the normal dataset-export command or Python API, then use the base-library ExportedDataset definition for the remote root.

For S3 exports, ExportedDataset loads config.yaml, metadata, and the small DALI .idx files. Tensor payloads remain in S3 and are read directly by DALI rather than downloaded through Python/boto3. Remote exports also avoid the per-sample Python WebDataset index; the tar/index lists in config.yaml are sufficient for DALI.

NexuML passes each s3://...tar path directly to the existing DALI WebDataset reader. DALI's WebDataset index parser requires local files, so NexuML stages only the corresponding .idx files in a worker-local temporary directory. DALI still owns sharding through shard_id=global_rank and num_shards=world_size.

Verify direct S3 tar support in the target DALI environment

DALI supports S3 in several readers, but current readers.webdataset documentation does not explicitly guarantee cloud URLs. NexuML therefore keeps this boundary deliberately thin: direct tar URLs are the intended path, and a real target-environment integration test must confirm them before production use. If that combination needs a fallback, add the smallest bounded adapter demonstrated by that test rather than a second cache/capability framework.

Credentials/endpoints come from the normal AWS/provider environment and supported S3 options, not from persisted secrets in the scenario.

Temporary KubeRay clusters

Cluster creation is infrastructure policy, not a second NexuML execution model. Use an infrastructure-owned KubeRay RayJob/cluster template to define images, CUDA, node selectors, queues, tolerations, service accounts, volumes, and autoscaling. Its entrypoint can remain the same nexuml train ... command.

Troubleshooting

  • Ray import error → install nexuml[ray] in the driver/worker environment.
  • Version mismatch → align Ray/Python versions between the project and cluster.
  • Post-train/evaluation rejection → current distributed global-finalization semantics are intentionally unsupported; run that scenario locally.
  • S3 failure → verify credentials, endpoint/region settings, and worker access to the S3 service.
  • DALI failure → install a compatible DALI build separately and verify python -c "import nvidia.dali" in the worker environment.