cp: Add SlurmRayClient for multi-node Ray cluster management on SLURM (1712) into r1.2.0

[svcnvidia-nemo-ci]

beep boop [🤖]: Hi @abhinavg4 👋,

we've cherry picked #1712 into  for you! 🚀

Please review and approve this cherry pick by your convenience!

[svcnvidia-nemo-ci]

/ok to test f061ff7

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This cherry-pick from #1712 adds SlurmRayClient, a RayClient subclass that handles multi-node Ray cluster management on SLURM by auto-detecting node roles (SLURM_NODEID), coordinating the head/worker startup via a shared port-broadcast file, and blocking worker processes until the cluster is torn down. It also adds tutorial scripts (pipeline.py, submit.sh, submit_container.sh, README.md) demonstrating the feature end-to-end.

Prior review threads have already flagged the two most significant issues (the 600 s hard-coded _read_head_port timeout not tied to worker_connect_timeout_s, and RAY_TMPDIR being exported but never consumed by SlurmRayClient). One additional usability concern remains in the tutorial submit scripts.

Confidence Score: 5/5

Safe to merge; all remaining findings are P2 style/usability suggestions that do not block correctness on the default happy path.

The core logic is well-implemented and comprehensively tested. Prior review threads already caught the two most impactful issues (hardcoded 600 s port-read timeout and the no-op RAY_TMPDIR export). The two new findings are both P2: a minor timing asymmetry in _wait_for_workers that only matters when worker_connect_timeout_s < 120 s (not the default), and a tutorial script usability gap where the logs/ directory must be pre-created before sbatch submission.

tutorials/slurm/submit.sh and submit_container.sh for the logs/ pre-creation note; nemo_curator/core/client.py _wait_for_workers for the deadline placement.

Important Files Changed

Filename	Overview
nemo_curator/core/client.py	Adds SlurmRayClient (head/worker role detection, port-broadcast via shared FS, worker connect timeout); well-structured but worker-join deadline starts before ray.init and _read_head_port uses a hardcoded 600 s timeout independent of worker_connect_timeout_s (flagged in prior review).
tests/core/test_slurm_ray_client.py	Comprehensive unit tests covering helper functions, head/worker lifecycle branches, port-file roundtrip, timeout cases, and cleanup; good use of monkeypatching with a fake Ray module.
tutorials/slurm/submit.sh	Working sbatch script; #SBATCH --output references a logs/ subdirectory that must be pre-created before submission — the in-script mkdir -p comes too late. RAY_TMPDIR is exported but not consumed by SlurmRayClient (flagged in prior review).
tutorials/slurm/submit_container.sh	Container-based sbatch script with same logs/ pre-creation requirement as submit.sh; otherwise correct Pyxis/enroot pattern.
tutorials/slurm/pipeline.py	Clean demo pipeline illustrating RayClient vs SlurmRayClient; CPU-only stages with correct try/finally cleanup.
tutorials/slurm/README.md	Good documentation of head/worker model and environment variables, but RAY_TMPDIR table entry implies the env var is read by SlurmRayClient when it is not (prior review thread); missing note about pre-creating logs/ directory.

Sequence Diagram

sequenceDiagram
    participant SLURM as SLURM (srun)
    participant Head as Node 0 (Head)
    participant Worker as Node N (Worker)
    participant FS as Shared FS (RAY_PORT_BROADCAST_DIR)

    SLURM->>Head: launch python pipeline.py --slurm
    SLURM->>Worker: launch python pipeline.py --slurm

    Head->>Head: _cleanup_local_ray() [if cleanup_on_start]
    Worker->>Worker: _cleanup_local_ray() [if cleanup_on_start]

    Head->>Head: super().start() → ray start --head
    Head->>Head: _write_head_port(job_id)
    Head->>FS: write ray_head_port_<job_id>

    Worker->>FS: _read_head_port(job_id) [polls until file appears]
    FS-->>Worker: port number

    Worker->>Worker: _run_as_worker(head_ip:port)
    Note over Worker: ray start --block --address=head:port (blocks until cluster torn down)

    Head->>Head: _wait_for_workers() [polls ray.nodes()]
    Note over Head: ray.init → polls until all N nodes alive

    Head->>Head: Pipeline runs (returns from start())
    Head->>Head: pipeline.run() completes
    Head->>Head: ray_client.stop() → ray stop + remove port file
    Note over Worker: ray process exits (cluster torn down)
    Worker->>SLURM: sys.exit(0)

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'r1.2.0' into cherry-pick-1..." | Re-trigger Greptile}

[greptile-apps]

Worker port-read timeout not tied to worker_connect_timeout_s

_read_head_port is called here without a timeout_s argument, so it always uses its hardcoded default of 600 s, regardless of self.worker_connect_timeout_s (default 300 s). A user who sets worker_connect_timeout_s=60 expecting fast failure on both sides will find workers silently waiting up to 600 s for the port file before timing out.

Suggested change

	actual_port = self._read_head_port(slurm_job_id)
	actual_port = self._read_head_port(slurm_job_id, timeout_s=self.worker_connect_timeout_s)

[greptile-apps]

RAY_TMPDIR export is not consumed by SlurmRayClient

SlurmRayClient (via RayClient) uses self.ray_temp_dir, which defaults to DEFAULT_RAY_TEMP_DIR (/tmp/ray). Neither _detect_slurm_resources nor any other method reads RAY_TMPDIR to set ray_temp_dir. The same pattern repeats in submit_container.sh line 75. As-is, Ray picks up the --temp-dir flag passed explicitly by the client, so this export is at best a no-op and may mislead users who expect it to control where Ray stores its files.

Consider passing the desired path as an explicit constructor argument instead:

SlurmRayClient(ray_temp_dir=f"/tmp/ray_{os.environ['SLURM_JOB_ID']}")

or remove the export RAY_TMPDIR lines if the intent was to rely on the client's default.

[abhinavg4]

/ok to test 47e9341