[Slurm] Multi-node clusters by kevinmingtarja · Pull Request #8219 · skypilot-org/skypilot

kevinmingtarja · 2025-12-06T09:52:20Z

This PR implements multi-node clusters support for Slurm.

Tested (run the relevant ones):

Code formatting: install pre-commit (auto-check on commit) or bash format.sh
Any manual or new tests for this PR (please specify below)
All smoke tests: /smoke-test (CI) or pytest tests/test_smoke.py (local)
Relevant individual tests: /smoke-test -k test_name (CI) or pytest tests/test_smoke.py::test_name (local)
- /smoke-test --slurm -k test_cli_logs
- /smoke-test --slurm -k test_launch_with_failing_setup
- /smoke-test --slurm -k test_multi_echo
- /smoke-test --slurm -k test_worker_provision_logs_streaming
- /smoke-test --slurm -k minimal
- /smoke-test --slurm -k ray_basic 🟢
- /smoke-test --slurm -k ray_train 🟢
- /smoke-test --slurm -k env_check 🟢
- /smoke-test --slurm -k test_multi_hostname 🟢
- /smoke-test --slurm -k test_multi_node_failure 🟢
- /smoke-test --slurm -k test_cancel_pytorch 🟢
Backward compatibility: /quicktest-core (CI) or pytest tests/smoke_tests/test_backward_compat.py (local)

gemini-code-assist · 2025-12-06T09:52:48Z

Summary of Changes

Hello @kevinmingtarja, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances SkyPilot's integration with Slurm by introducing full support for multi-node clusters. It achieves this through a dedicated Slurm executor script that manages per-node task execution, environment variable propagation, and synchronized setup/run phases. The changes also involve refactoring the core task code generation to accommodate distributed logging and robust command execution, alongside updates to the Slurm cloud backend and test suite to validate the new capabilities.

Highlights

Slurm Multi-Node Support: Fully enabled multi-node cluster capabilities for Slurm, removing previous limitations and allowing distributed workloads.
Dedicated Slurm Executor: Introduced sky/skylet/executor/slurm.py, a new Python script to manage distributed task execution, environment variables, and logging on individual Slurm compute nodes.
Refactored Task Code Generation: Centralized log path handling and introduced a new build_task_runner_cmd function for robust srun command construction and script execution, including synchronization mechanisms via signal files.
Enhanced Setup and Logging: Improved the setup process for multi-node Slurm jobs, ensuring directory creation only on the head node and providing distinct per-node logging for setup and run tasks.
Updated Test Suite: Modified smoke and unit tests to reflect and validate the new multi-node Slurm functionality, removing previous Slurm-specific single-node restrictions.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

kevinmingtarja · 2025-12-06T09:58:47Z

/smoke-test --slurm -k test_cli_logs
/smoke-test --slurm -k test_launch_with_failing_setup
/smoke-test --slurm -k test_multi_echo
/smoke-test --slurm -k test_worker_provision_logs_streaming
/smoke-test --slurm -k minimal

gemini-code-assist

Code Review

This pull request introduces multi-node cluster support for Slurm, which is a significant enhancement. The changes involve a major refactoring of the Slurm code generation logic, moving it to a dedicated remote executor module (sky.skylet.executor.slurm). This is a good architectural decision that improves modularity and maintainability. The implementation details for multi-node support, such as node IP discovery, per-node logging, and synchronization, appear to be well-thought-out. However, I've identified a critical issue in the provisioning script that could prevent clusters from being set up correctly. Overall, this is a great step forward, and with the identified issue fixed, it will be a solid contribution.

kevinmingtarja · 2025-12-06T10:07:03Z

/smoke-test --slurm -k test_multi_echo

kevinmingtarja · 2025-12-06T10:08:40Z

/smoke-test --slurm -k test_multi_echo

kevinmingtarja · 2025-12-07T06:14:49Z

/smoke-test --slurm -k ray_basic
/smoke-test --slurm -k env_check
/smoke-test --slurm -k test_multi_hostname
/smoke-test --slurm -k test_multi_node_failure
/smoke-test --slurm -k test_cancel_pytorch

kevinmingtarja · 2025-12-07T06:28:17Z

/smoke-test --slurm -k test_cancel_pytorch

kevinmingtarja · 2025-12-07T06:50:44Z

/smoke-test --slurm -k test_cancel_pytorch
/smoke-test --slurm -k env_check

kevinmingtarja · 2025-12-07T06:55:08Z

/smoke-test --slurm -k ray_train

kevinmingtarja · 2025-12-07T08:13:20Z

/smoke-test --slurm -k test_multi_node_failure

Michaelvll · 2025-12-08T05:31:51Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

kevinmingtarja · 2025-12-08T19:22:37Z

/smoke-test --slurm -k test_cancel_pytorch

SeungjinYang · 2025-12-08T19:59:07Z

            'has_setup_cmd = False',
            'setup_cmd = None',
            'setup_envs = None',
-            'setup_log_path = None',
+            'setup_log_dir = None',
            'setup_num_nodes = None',
        ]


I'm generally a fan of moving things out of the string scripts for maintainability, since any scripts in string doesn't benefit from automated tools, etc. What do you think about doing something here like:

self._code += [ 'autostop_lib.set_last_active_time_to_now()', f'job_lib.set_status({job_id!r}, job_lib.JobStatus.PENDING)', ] self._has_setup_cmd = False self._setup_cmd = None, self._setup_envs = None, self._setup_log_dir = None, self._setup_num_nodes = None,

and build the command dynamically on add_task?

Sounds good, that would indeed be cleaner.

Nice! I have more ideas on how to clean up this part more, but let's merge this PR in and I can make a separate PR incorporating those cleanups. The code is functional, and that's what matters for now.

…t_runtime_dir cleanup

SeungjinYang · 2025-12-08T23:40:56Z

+            # Clean up sky runtime directory on each node.
+            srun --nodes={num_nodes} rm -rf {skypilot_runtime_dir}
            rm -rf {sky_home_dir}
        }}
        trap cleanup TERM

-        # Create sky directory for the cluster.
-        mkdir -p {sky_home_dir} {skypilot_runtime_dir}
+        # Create sky home directory for the cluster.
+        mkdir -p {sky_home_dir}
+        # Create sky runtime directory on each node.
+        srun --nodes={num_nodes} mkdir -p {skypilot_runtime_dir}


When you do multiple sruns of multiple nodes, are these sruns guarenteed to run on the same set of nodes?

In this case yes, because it is being run from inside an sbatch script.

But actually in this case, it's just because we do srun --nodes N (N number of nodes requested by sbatch).

If we did srun --nodes M, M < N, I think there are no guarantees that multiple sruns would run on the same subset of M nodes.

That makes sense, it is a quirk so I'd ideally like it jotted down in a comment somewhere

SeungjinYang · 2025-12-08T23:41:24Z

-        'failed': status_lib.ClusterStatus.INIT,
+        'failed': None,


I'm curious what this change was motivated by

Ah yes so apparently when we sky down a cluster in Slurm, under the hood we will call scancel, and the job will go into failed state and will remain in squeue for a few minutes:

$ squeue --states failed JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON) 3111 CPUs sky-93c7 ubuntu F 40:50 2 (NonZeroExitCode)

And so then if we relaunch a new cluster with the same name, our status refresh / query_instances would include the status of this previously downed cluster in the statuses map it returns, which would be confusing.

So by making it None, it is omitted when non_terminated_only is True.

Ah ok, interesting. Could we jot this done in a comment somewhere as well

SeungjinYang

Thanks @kevinmingtarja, awesome work! This PR is ready to be merged in my opinion, once all the tests listed in the PR description is run successfully.

kevinmingtarja · 2025-12-09T02:07:00Z

/smoke-test --slurm -k ray_train
/smoke-test --slurm -k env_check
/smoke-test --slurm -k test_multi_hostname
/smoke-test --slurm -k test_multi_node_failure
/smoke-test --slurm -k test_cancel_pytorch

kevinmingtarja · 2025-12-09T02:24:00Z

/smoke-test --slurm -k test_cli_logs
/smoke-test --slurm -k test_launch_with_failing_setup
/smoke-test --slurm -k test_worker_provision_logs_streaming
/smoke-test --slurm -k minimal

kevinmingtarja · 2025-12-09T02:49:34Z

/smoke-test -k test_job_pipeline --slurm --jobs-consolidation

kevinmingtarja · 2025-12-09T02:50:07Z

/quicktest-core

kevinmingtarja · 2025-12-09T03:02:53Z

/smoke-test --slurm -k ray_basic

kevinmingtarja added 5 commits December 6, 2025 01:37

[Slurm] Multi-node clusters

4ef9300

update snapshot test

4fc24c7

re-enable disabled smoke tests

407fdfc

properly compute SKYPILOT_NODE_IPS

43d293e

update snapshot test

b8d8bc1

fix test_multi_echo for slurm

a89f766

gemini-code-assist Bot reviewed Dec 6, 2025

View reviewed changes

Comment thread sky/provision/slurm/instance.py Outdated

no spot

bfd7db9

kevinmingtarja added 4 commits December 6, 2025 22:09

fix log prefix for multi node

c1741ec

re-enable disabled smoke tests

a07f6c9

add copy_skypilot_templates_commands to slurm yml

5673b7b

fix gpu check in test_examples

6e44477

kevinmingtarja added 2 commits December 6, 2025 22:26

update snapshot test

4971d15

fix gpu check in test_cancel_pytorch

49100c7

kevinmingtarja added 2 commits December 6, 2025 22:36

fix get_available_gpus

08cc897

fix setup no run

215416f

fix FAILED_SETUP handling

4ed03e9

chatgpt-codex-connector Bot reviewed Dec 8, 2025

View reviewed changes

Comment thread sky/provision/slurm/instance.py Outdated

fix test_cancel_pytorch

d1a7129

SeungjinYang reviewed Dec 8, 2025

View reviewed changes

Comment thread sky/provision/slurm/instance.py Outdated

Comment thread sky/backends/task_codegen.py

Comment thread sky/backends/task_codegen.py Outdated

Comment thread sky/backends/task_codegen.py Outdated

SeungjinYang reviewed Dec 8, 2025

View reviewed changes

Comment thread sky/backends/task_codegen.py

docs on task executor architecture

90ae73a

SeungjinYang reviewed Dec 8, 2025

View reviewed changes

kevinmingtarja added 2 commits December 8, 2025 13:16

wait for sbatch script to finish setup before proceeding

4dc9b06

Make query_instances more robust against terminated jobs, fix skypilo…

2cfa9e2

…t_runtime_dir cleanup

SeungjinYang reviewed Dec 8, 2025

View reviewed changes

kevinmingtarja added 4 commits December 8, 2025 16:36

clean up SlurmCodeGen

18ea22c

dont pass in cluster_num_nodes

6734614

define SKY_SLURM_PYTHON_CMD

64d30b9

add more todos/comments around NFS assumptions we're making

d2de278

kevinmingtarja marked this pull request as ready for review December 9, 2025 00:59

kevinmingtarja added 3 commits December 8, 2025 17:09

update snapshot test

dce9b4b

use which

c982133

add more comments

d6ceba0

SeungjinYang approved these changes Dec 9, 2025

View reviewed changes

kevinmingtarja merged commit fced388 into master Dec 9, 2025
30 of 32 checks passed

kevinmingtarja deleted the slurm-multi-node branch December 9, 2025 04:10

Uh oh!

Conversation

kevinmingtarja commented Dec 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot commented Dec 6, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

kevinmingtarja commented Dec 6, 2025

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

kevinmingtarja commented Dec 6, 2025

Uh oh!

kevinmingtarja commented Dec 6, 2025

Uh oh!

kevinmingtarja commented Dec 7, 2025

Uh oh!

kevinmingtarja commented Dec 7, 2025

Uh oh!

kevinmingtarja commented Dec 7, 2025

Uh oh!

kevinmingtarja commented Dec 7, 2025

Uh oh!

kevinmingtarja commented Dec 7, 2025

Uh oh!

Michaelvll commented Dec 8, 2025

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

kevinmingtarja commented Dec 8, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinmingtarja Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinmingtarja Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SeungjinYang left a comment

Choose a reason for hiding this comment

Uh oh!

kevinmingtarja commented Dec 9, 2025

Uh oh!

kevinmingtarja commented Dec 9, 2025

Uh oh!

kevinmingtarja commented Dec 9, 2025

Uh oh!

kevinmingtarja commented Dec 9, 2025

Uh oh!

kevinmingtarja commented Dec 6, 2025 •

edited

Loading

kevinmingtarja Dec 9, 2025 •

edited

Loading

kevinmingtarja Dec 9, 2025 •

edited

Loading