revert: "chore: improve ray.sub generalization across clusters"

[terrykong]

Reverts #1451

We discovered on our cluster that while each ray node claims all the cpus, for some reason the CPU affinity of a ray task will not be all the CPUS and only two of them. This reverts this problematic change until we can figure

Summary by CodeRabbit

• New Features

  • Automatic GPU resource detection with validation and clear logs
  • Dynamic CPU-per-worker defaulting based on GPUs and applied across cluster

• Bug Fixes

  • Consistent task and CPU allocation for head and worker processes
  • Improved node IP discovery with multiple resolution methods and mapping

• Improvements

  • Unified resource argument handling for all cluster commands and scripts
  • More robust cluster startup and validation behavior

[coderabbitai]

📝 Walkthrough

Walkthrough

Detects SLURM GRES support via a new maybe_gres_arg, composes COMMON_SRUN_ARGS dynamically (including GRES, container mounts, workdir), replaces static srun flags with --ntasks=1 and --cpus-per-task (CPUS_PER_WORKER = GPUS_PER_NODE*16), updates all srun uses, and improves node IP resolution with fallbacks and an ip_addresses_array.

Changes

Cohort / File(s)	Summary
GRES detection & init ray.sub	Adds maybe_gres_arg to detect SLURM GRES/gpu support, validate GPUS_PER_NODE against detected GRES GPUs, set GRES_ARG, and log presence/absence.
COMMON srun args composition ray.sub	Replaces static COMMON_SRUN_ARGS with dynamic composition including GRES_ARG, container settings, mounts, container workdir; derives CPUS_PER_WORKER = GPUS_PER_NODE * 16.
srun invocation updates ray.sub	Changes head/worker/utility srun calls to include --ntasks=1, --cpus-per-task=$CPUS_PER_WORKER; removes --exclusive, adds --exact for workers, and ensures GRES_ARG is applied consistently.
Attach / driver scripts ray.sub	Ensures generated attach/driver commands include GRES_ARG, --ntasks=1, and --cpus-per-task for head and workers.
IP resolution & mapping ray.sub	Enhances node IP resolution with fallbacks (host, getent, nslookup, ping), detailed logging, and builds ip_addresses_array mapping for head and workers.

Sequence Diagram(s)

sequenceDiagram
    participant Init as startup
    participant GRES as maybe_gres_arg
    participant Composer as COMMON_SRUN_ARGS composer
    participant Head as srun head
    participant Worker as srun worker
    participant Resolver as IP resolver

    Init->>GRES: detect GRES (SLURM_NODE_GRES/GRES config)
    GRES-->>Init: GRES_ARG or empty (validate GPUS_PER_NODE)
    Init->>Composer: compose COMMON_SRUN_ARGS (include GRES_ARG, mounts, workdir, CPUS_PER_WORKER)
    Composer-->>Head: head srun command (--ntasks=1 --cpus-per-task, GRES_ARG)
    Composer-->>Worker: worker srun command (--ntasks=1 --cpus-per-task, --exact, GRES_ARG)
    Head->>Resolver: resolve head IP (host/getent/nslookup/ping)
    Worker->>Resolver: resolve worker IPs (build ip_addresses_array)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

• Review maybe_gres_arg detection and mismatch/error behavior between detected GRES and GPUS_PER_NODE.
• Verify all srun call sites were updated to use --ntasks=1 and correct --cpus-per-task values (head, worker, monitoring, status checks).
• Validate CPUS_PER_WORKER derivation and any assumptions (GPUS_PER_NODE presence/zero).
• Confirm IP resolution fallbacks produce stable, consistent node-to-IP mappings in target environments.

Possibly related PRs

• chore: improve ray.sub generalization across clusters #1451 — Modifies SLURM-related logic in ray.sub (GRES detection, COMMON_SRUN_ARGS, CPU/GPU allocation); directly related and may conflict.

Suggested reviewers

• hemildesai
• parthmannan

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR is a major revert of ray.sub changes (SLURM GRES, GPU detection, CPU allocation) with no test results or testing documentation in the description.	Add testing documentation including reason for revert, test results, validation that no regressions occur, and evidence the reverted code works correctly.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title clearly indicates this is a revert of a previous change (#1451), which is directly reflected in the raw_summary showing removal of GRES support and dynamic SLURM argument composition that were added in that change.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch revert-1451-tk/ray-sub-lambda

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6a035bc and 8e21f2a.

📒 Files selected for processing (1)

• ray.sub (9 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: Docs_Tests
• GitHub Check: Lint check
• GitHub Check: Post submodule check comment / Comment on PR
• GitHub Check: Post automodel integration comment / Comment on PR