Add support for notebooks

[jacobtomlinson]

I just tested out the quickstart on Google Colab but it just hangs.

• Sign up for Google Colab
• Choose a T4 Runtime type
• Run !pip install nsight-python
• Run the quickstart code

import torch
import nsight

@nsight.analyze.kernel
def benchmark_matmul(n):
    """
    The simplest possible benchmark.
    We create two matrices and multiply them.
    """
    # Create two NxN matrices on GPU
    a = torch.randn(n, n, device="cuda")
    b = torch.randn(n, n, device="cuda")

    # Mark the kernel we want to profile
    with nsight.annotate("matmul"):
        c = a @ b

    return c

result = benchmark_matmul(1024)

After running for 10 mins it is still executing.

[jacobtomlinson]

I also reproduced this in a fresh environment locally in Jupyter.

mkdir /tmp/nsight-test && cd /tmp/nsight-test 
pixi init
pixi shell
pixi add pytorch-gpu jupyterlab
pixi add --pypi nsight-python
pixi run jupyter lab

Then run the quickstart example in a notebook and it just hangs.

# pixi.toml
[workspace]
channels = ["conda-forge"]
name = "nsight-test"
platforms = ["linux-64"]
version = "0.1.0"

[tasks]

[dependencies]
jupyterlab = ">=4.5.0,<5"
pytorch-gpu = ">=1.11.0,<2"

[pypi-dependencies]
nsight-python = ">=0.9.4, <0.10"

[Alok-Joshi]

The issue occurs because nsight-python profiles decorated functions by relaunching the entire Python script that contains them under NVIDIA Nsight Compute.

This fails in Colab because each code cell execution launches /usr/local/lib/python3.12/dist-packages/colab_kernel_launcher.py, which then starts an IPython kernel. nsight-python mistakenly treats this launcher script as the one containing the decorated function and tries to run it under Nsight Compute, causing the process to hang. We’re actively exploring a better long-term solution. In the meantime, you can use the following workaround:

%%writefile quickstart.py
import torch
import nsight

@nsight.analyze.kernel
def benchmark_matmul(n):
    """
    The simplest possible benchmark.
    We create two matrices and multiply them.
    """
    # Create two NxN matrices on GPU
    a = torch.randn(n, n, device="cuda")
    b = torch.randn(n, n, device="cuda")

    # Mark the kernel we want to profile
    with nsight.annotate("matmul"):
        c = a @ b

    return c

result = benchmark_matmul(1024)

Followed by :

%run quickstart.py

[jacobtomlinson]

That makes sense. Both times I tried to use it were in notebooks, Colab and Jupyter. I expect this is a common place for people to use this tool.

Thank you for the workaround, I managed to get things running on Colab.

I was also interested in plotting the results, which get saved to a file called plot.png by default. With this workaround I also needed to use IPython.display.Image to view the plot.

!pip install nsight-python

%%writefile quickstart.py
import torch
import nsight

@nsight.analyze.plot()
@nsight.analyze.kernel()
def benchmark_matmul(n):
    """
    The simplest possible benchmark.
    We create two matrices and multiply them.
    """
    # Create two NxN matrices on GPU
    a = torch.randn(n, n, device="cuda")
    b = torch.randn(n, n, device="cuda")

    # Mark the kernel we want to profile
    with nsight.annotate("matmul"):
        c = a @ b

    return c

result = benchmark_matmul(configs=[(1024, ), (2048, ), (4096, )])

%run quickstart.py

from IPython.display import Image
Image("plot.png")

[Alok-Joshi]

We do not have a complete fix for this yet. Will document this in known issues.