CUDA, Python and JupyterLab

echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
nvcc --version | tail -n 2          # expect: Cuda compilation tools, release 12.6, V12.6.68
ls -l /usr/local/cuda               # points to /usr/local/cuda-12.6

cd ~/sutraai
mkdir -p notebooks sutra_core configs data/raw data/processed models exports/golden reports
touch sutra_core/__init__.py
python3 --version                                   # expect: Python 3.10.12
python3 -m venv .venv                               # isolated: Ubuntu's old Python packages stay out
cat > constraints-jp621.txt <<'EOF'
# SutraAI tested set: Jetson Orin NX 16GB · JetPack 6.2.1 · CUDA 12.6.68 · cuDNN 9.3.0 · TensorRT 10.3.0.30 · Python 3.10.12
torch==2.8.0
torchvision==0.23.0
torchaudio==2.8.0
onnxruntime-gpu==1.23.0
numpy==1.26.4
matplotlib==3.10.9
pillow==11.3.0
scipy==1.15.3
EOF
echo 'export PIP_CONSTRAINT=$HOME/sutraai/constraints-jp621.txt' >> .venv/bin/activate   # every pip install in this venv obeys the lock
source .venv/bin/activate
pip install --upgrade pip
pip install numpy==1.26.4

# Link only JetPack's TensorRT 10.3 Python bindings into the venv (nothing else from Ubuntu)
TRT_DIR=$(/usr/bin/python3 -c "import tensorrt, os; print(os.path.dirname(os.path.dirname(tensorrt.__file__)))")
echo "$TRT_DIR"                                     # expect: /usr/lib/python3.10/dist-packages
SP="$VIRTUAL_ENV/lib/python3.10/site-packages"
for p in "$TRT_DIR"/tensorrt*; do ln -sfn "$p" "$SP/$(basename "$p")"; done
python -c "import tensorrt; print(tensorrt.__version__)"                          # expect: 10.3.0.30
python -c "import sys; print([p for p in sys.path if 'dist-packages' in p])"      # expect: []

cd ~/sutraai && source .venv/bin/activate
SP="$VIRTUAL_ENV/lib/python3.10/site-packages"

# 1) stop seeing Ubuntu's Python packages
sed -i 's/^include-system-site-packages = .*/include-system-site-packages = false/' .venv/pyvenv.cfg
grep include-system .venv/pyvenv.cfg                 # expect: include-system-site-packages = false

# 2) remove the two mpl_toolkits workaround files (no longer needed)
rm -f "$SP/00-sutraai-mpl-toolkits.pth" "$SP/mpl_toolkits/__init__.py"

# 3) link only TensorRT
TRT_DIR=$(/usr/bin/python3 -c "import tensorrt, os; print(os.path.dirname(os.path.dirname(tensorrt.__file__)))")
for p in "$TRT_DIR"/tensorrt*; do ln -sfn "$p" "$SP/$(basename "$p")"; done

# 4) install what pip skipped earlier because Ubuntu "already had it" (installed packages are kept)
pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 onnxruntime-gpu==1.23.0 --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/
pip install jupyterlab ipykernel ipywidgets matplotlib pandas scikit-learn tqdm pyyaml \
            timm onnx librosa soundfile openai sentence-transformers thop nbstripout ai-edge-litert
pip check                                            # expect: No broken requirements found.

# 5) verify
python -c "import sys; print([p for p in sys.path if 'dist-packages' in p])"      # expect: []
python -c "import torch, tensorrt, matplotlib; print(torch.__version__, torch.cuda.is_available(), tensorrt.__version__, matplotlib.__version__)"
# expect: 2.8.0 True 10.3.0.30 3.10.9

# a) OpenBLAS + OpenMP from Ubuntu
sudo apt-get update
sudo apt-get install -y libopenblas-dev libomp-dev

# b) cuDSS 0.7.1 for Jetson (CUDA 12)  →  provides libcudss.so.0
mkdir -p ~/downloads && cd ~/downloads
wget https://developer.download.nvidia.com/compute/cudss/redist/libcudss/linux-aarch64/libcudss-linux-aarch64-0.7.1.4_cuda12-archive.tar.xz
echo "f283c31b6badf4a5277014295705eac8e6e28f27de30d746719ea1cef7a750b8  libcudss-linux-aarch64-0.7.1.4_cuda12-archive.tar.xz" | sha256sum -c   # expect: OK
sudo mkdir -p /opt/nvidia
sudo tar xf libcudss-linux-aarch64-0.7.1.4_cuda12-archive.tar.xz -C /opt/nvidia
dirname "$(find /opt/nvidia/libcudss-linux-aarch64-0.7.1.4_cuda12-archive -name 'libcudss.so.0' | head -n1)" | sudo tee /etc/ld.so.conf.d/cudss.conf

# c) cuSPARSELt 0.6.3 for Jetson (the version PyTorch uses with CUDA 12.6)  →  provides libcusparseLt.so.0
wget https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-aarch64/libcusparse_lt-linux-aarch64-0.6.3.2-archive.tar.xz
echo "91501d0c05d1ff0dd67399ecd7c1bf76a620e842dce54ae4c8a1f07cec0673e5  libcusparse_lt-linux-aarch64-0.6.3.2-archive.tar.xz" | sha256sum -c   # expect: OK
sudo tar xf libcusparse_lt-linux-aarch64-0.6.3.2-archive.tar.xz -C /opt/nvidia
dirname "$(find /opt/nvidia/libcusparse_lt-linux-aarch64-0.6.3.2-archive -name 'libcusparseLt.so.0' | head -n1)" | sudo tee /etc/ld.so.conf.d/cusparselt.conf

# d) register both with the system loader
sudo ldconfig
ldconfig -p | grep -E "libcudss.so.0|libcusparseLt.so.0|libopenblas.so.0"   # expect 3+ lines
cd ~/sutraai

cd ~/sutraai && source .venv/bin/activate
pip install --no-cache-dir torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/
pip install --no-cache-dir onnxruntime-gpu==1.23.0 --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/

# check that every library torch links against is found (runs without importing torch)
TORCH_LIB=$(python -c "import importlib.util, pathlib; print(pathlib.Path(importlib.util.find_spec('torch').origin).parent / 'lib')")
for f in "$TORCH_LIB"/*.so*; do ldd "$f"; done | grep "not found" | sort -u          # expect: no output

python -c "import torch, torchvision, numpy; print(torch.__version__, torchvision.__version__, numpy.__version__, torch.version.cuda, torch.cuda.is_available())"
# expect: 2.8.0 0.23.0 1.26.4 12.6 True

pip install jupyterlab ipykernel ipywidgets matplotlib pandas scikit-learn tqdm pyyaml \
            timm onnx librosa soundfile openai sentence-transformers thop nbstripout ai-edge-litert
python -m ipykernel install --user --name sutraai --display-name "SutraAI (venv · torch 2.8 · CUDA 12.6)"
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"     # still 2.8.0 True
pip check                                                                        # expect: No broken requirements found.

# isolation check: every package the notebooks import must come from the venv
python - <<'EOF'
import importlib
for name in ["numpy", "matplotlib", "mpl_toolkits.mplot3d", "PIL", "scipy", "sklearn", "pandas", "torch", "torchvision", "onnxruntime", "tensorrt"]:
    try:
        m = importlib.import_module(name)
    except Exception as e:
        print(f"{name:22} FAILED: {type(e).__name__}: {e}"); continue
    where = "venv" if "/.venv/" in (m.__file__ or "") else "OUTSIDE THE VENV  <- check .venv/pyvenv.cfg (substep 2)"
    print(f"{name:22} {str(getattr(m, '__version__', '')):12} {where}")
EOF
# expect: every line ends with "venv" (tensorrt too: it's linked in)

source ~/sutraai/.venv/bin/activate
hf auth login        # paste the token when asked; if asked "Add token as git credential?", answer n
hf auth whoami       # expect: your Hugging Face username

tmux new -s lab                    # detach with Ctrl+B then D; come back with: tmux attach -t lab
cd ~/sutraai && source .venv/bin/activate
jupyter lab --no-browser --port 8888

ssh -N -L 8888:localhost:8888 <your-user>@<jetson-ip>      # or your Tailscale address when away from home
# then open http://localhost:8888 and paste the token printed in the Jetson terminal

import platform, numpy, torch, torchvision, onnxruntime as ort, matplotlib
print("Python     ", platform.python_version())                       # 3.10.12
print("numpy      ", numpy.__version__)                               # 1.26.4
print("matplotlib ", matplotlib.__version__, matplotlib.__file__)     # 3.10.9 …/.venv/lib/python3.10/site-packages/…
print("torch      ", torch.__version__, "| CUDA available:", torch.cuda.is_available())   # 2.8.0 True
print("torchvision", torchvision.__version__)                         # 0.23.0
print("CUDA (torch build):", torch.version.cuda)                      # 12.6
print("cuDNN      ", torch.backends.cudnn.version())                  # 9.3.x reported as 903xx
print("GPU        ", torch.cuda.get_device_name(0))                   # Orin
free, total = torch.cuda.mem_get_info()
print(f"GPU-visible memory: {free/1e9:.1f} GB free of {total/1e9:.1f} GB")   # about 15 GB total on Orin NX 16GB
try:
    import tensorrt
    print("TensorRT   ", tensorrt.__version__)                         # 10.3.0.30
except ImportError:
    print("TensorRT not visible: run the TensorRT link lines in substep 2")
print("ONNX Runtime", ort.__version__, ort.get_available_providers())  # 1.23.0, includes CUDAExecutionProvider

import torch, time

def bench(device, n=2048, reps=10, dtype=torch.float32):
    a = torch.randn(n, n, device=device, dtype=dtype)
    b = torch.randn(n, n, device=device, dtype=dtype)
    for _ in range(2): a @ b                          # warm-up
    if device == "cuda": torch.cuda.synchronize()
    t = time.perf_counter()
    for _ in range(reps): a @ b
    if device == "cuda": torch.cuda.synchronize()
    return (time.perf_counter() - t) / reps * 1000    # ms per multiply

cpu, gpu, gpu16 = bench("cpu"), bench("cuda"), bench("cuda", dtype=torch.float16)
print(f"CPU (8 cores) {cpu:.1f} ms | GPU fp32 {gpu:.1f} ms | GPU fp16 {gpu16:.1f} ms")

cd ~/sutraai && git init
nbstripout --install                       # strips notebook outputs on commit
printf ".venv/\n.venv-convert/\ndata/\nmodels/\nexports/\n*.engine\n" > .gitignore
pip freeze > requirements-lock.txt         # full snapshot of what actually got installed
git add . && git commit -m "Step 02: environment ready (JetPack 6.2.1, CUDA 12.6.68, torch 2.8.0)"