#!/usr/bin/env bash
set -euo pipefail

if ! command -v nvidia-ctk >/dev/null 2>&1; then
  echo "nvidia-ctk is not installed. Install NVIDIA Container Toolkit first." >&2
  exit 1
fi

sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json

if command -v systemctl >/dev/null 2>&1; then
  sudo systemctl restart docker
else
  sudo service docker restart
fi

docker info --format '{{json .Runtimes}}'
docker run --rm --runtime nvidia --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi