Base Python
A clean Python environment for interactive analysis and everyday computation.
Interactive computing for Jupyter
Super Kernel
Bring supercomputing into your daily notebook.
Keep the familiar notebook workflow while choosing the computing runtime that fits your work—from everyday Python to MPI and PyTorch.
Jupyter Notebook
→
Super Kernel
→
Your Runtime
One notebook, the right runtime
Choose a runtime in JupyterLab and work with ordinary notebook cells.
MPI
Run the same notebook cell across multiple processes and computing nodes.
PyTorch
Use a ready-to-work PyTorch environment for model and tensor workloads.
Distributed Python with MPI
Select an MPI Super Kernel and use ordinary Python. No special cell magic is required.
In [1]:
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
print(f"Hello from rank {rank} of {size}")
Out:
Hello from rank 0 of 4 Hello from rank 1 of 4 Hello from rank 2 of 4 Hello from rank 3 of 4
Distributed tensors with PyTorch
Run one ordinary notebook cell across every PyTorch process, using Gloo on CPU or NCCL when GPUs are available.
In [1]:
import os
import torch
import torch.distributed as dist
local_rank = int(os.environ["LOCAL_RANK"])
use_cuda = torch.cuda.is_available()
device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
if use_cuda:
torch.cuda.set_device(device)
if not dist.is_initialized():
dist.init_process_group(backend="nccl" if use_cuda else "gloo")
rank = dist.get_rank()
value = torch.tensor(float(rank + 1), device=device)
dist.all_reduce(value, op=dist.ReduceOp.SUM)
print(f"Rank {rank}/{dist.get_world_size()}: sum = {value.item()}")
Out:
Rank 0/4: sum = 10.0 Rank 1/4: sum = 10.0 Rank 2/4: sum = 10.0 Rank 3/4: sum = 10.0