Skip to content

Slurm notebook - nccl test mappings to GPU skus #64

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
radu-malliu wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Slurm notebook - nccl test mappings to GPU skus #64

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2e818f3
add instance type to nccl test mapping; add instance type selector
radu-malliu Apr 3, 2026
af7fdf4
fix: remove dependencies
radu-malliu Apr 3, 2026
b393e92
fix: use ui components
radu-malliu Apr 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 66 additions & 35 deletions ...na_sunk_cluster_walkthrough_benchmarks.py → arena/notebooks/sunk_cluster_walkthrough.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,14 @@
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "k8s==0.28.0",
# "kubernetes==35.0.0",
# "marimo>=0.20.2",
# "ruamel-yaml>=0.19.1",
# "typing-extensions>=4.15.0"
# ]
# ///

import marimo

__generated_with = "0.19.7"
Expand All @@ -17,31 +28,22 @@ def _():
# setup cell runs before anything else, recommend putting values in here
import marimo as mo

"""
Import SSH helpers from arena library.

Configuration via environment variables:
CW_ARENA_SSH_KEY_PATH: Path to SSH private key (default: /root/.ssh/id_rsa)
CW_ARENA_SSH_HOST: SSH host (e.g., user+tenant@sunk.tenant.coreweave.app)
"""
from lib.remote_execution_helpers import shell

return mo, shell


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
# CoreWeave AI Labs: SUNK Cluster

/// admonition | About This Notebook
type: info

This notebook provides a walkthrough for inspecting and benchmarking your SUNK (Slurm on Kubernetes) cluster.
///
""")
return

from lib.k8s import K8s

TEST_NAMES = {

@MYanello MYanello Apr 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider populating this with what's in the test folder

"gb300-4x": "nccl-test-distributed-gb200-nvl72-enroot.slurm",
"gb300-4x-e": "nccl-test-distributed-gb300-roce-nvl72-enroot.slurm",
"gb200-4x": "nccl-test-distributed-gb200-nvl72-enroot.slurm",
"b300-8x": "nccl-test-distributed-h100-64.slurm",
"b200-8x": "nccl-test-distributed-h100-64.slurm",
"gd-8xh200ib-i128":"nccl-test-distributed-h100-64.slurm",
"gd-8xh100ib-i128":"nccl-test-distributed-h100-64.slurm",
"gd-8xa100-i128": "nccl-test-distributed-a100-64.slurm",
# L40S, L40, GH200, RTX Pro 6000 — no matching nccl test script
}

return mo, shell, K8s, TEST_NAMES

@app.cell(hide_code=True)
def _(mo):
Expand Down Expand Up @@ -84,21 +86,21 @@ def _(mo):
@app.cell
def _(shell):
# Slurm cluster info:
print(shell("sinfo"))
shell("sinfo")
return


@app.cell
def _(shell):
# Slurm user info:
print(shell("sacctmgr show users"))
shell("sacctmgr show users")
return


@app.cell
def _(shell):
# Slurm accounting info:
print(shell("sacctmgr show associations format=User,Account,Partition,QOS"))
shell("sacctmgr show associations format=User,Account,Partition,QOS")
return


Expand All @@ -113,7 +115,7 @@ def _(mo):
@app.cell
def _(shell):
# Node info
print(shell("scontrol show nodes"))
shell("scontrol show nodes")
return


Expand All @@ -127,7 +129,7 @@ def _(mo):

@app.cell
def _(shell):
print(shell("scontrol show partition"))
shell("scontrol show partition")
return


Expand Down Expand Up @@ -172,6 +174,24 @@ def _(mo, num_nodes):
return


@app.cell(hide_code=True)
def _(mo, K8s):
k8s = K8s()
nodes = k8s.nodes
gpu_nodes = nodes.get("gpu") or {}
gpu_keys = list(gpu_nodes.keys())
if not gpu_keys:
gpu_keys = ["(no GPU nodes detected)"]
default_node_type = gpu_keys[0]
node_type_dropdown = mo.ui.dropdown(
options=gpu_keys, value=default_node_type, label="Node type"
)
_gpu_type_ui = mo.md(
f"**GPU node types (from cluster API)**\n\n{node_type_dropdown}"
)
_gpu_type_ui
return (node_type_dropdown,)

@app.cell(hide_code=True)
def _(mo):
submit_btn = mo.ui.run_button(label="Submit NCCL Test Job")
Expand All @@ -180,25 +200,36 @@ def _(mo):


@app.cell
def _(num_nodes, shell, submit_btn):
def _(
mo,
num_nodes,
shell,
submit_btn,
node_type_dropdown: mo.ui.dropdown,
TEST_NAMES,
):
if submit_btn.value:
cmd = f"cd /mnt/data/arena/benchmarks/nccl/nccl-tests/slurm && sbatch -N {num_nodes.value} nccl-test-distributed-h100-64.slurm"
print(f"Running: sbatch -N {num_nodes.value} ...")
print(shell(cmd))
if node_type_dropdown.value not in TEST_NAMES:
print(f"No NCCL slurm script mapped for node type {node_type_dropdown.value!r}.")
else:
script = TEST_NAMES[node_type_dropdown.value]
cmd = f"cd /mnt/data/arena/benchmarks/nccl/nccl-tests/slurm && sbatch -N {num_nodes.value} {script}"
print(f"Running: sbatch -N {num_nodes.value} ...")
shell(cmd)
else:
print("Click 'Submit NCCL Test Job' to run the benchmark")
return


@app.cell
def _(shell):
print(shell("squeue"))
shell("squeue")
return


@app.cell
def _(shell):
print(shell("ls /mnt/data/arena/benchmarks/nccl/nccl-tests/slurm/*.out"))
shell("ls /mnt/data/arena/benchmarks/nccl/nccl-tests/slurm/*.out")
return


Expand Down