#!/usr/bin/env bash
# Purpose: pull the NVIDIA PyTorch container on a DGX Spark, check that PyTorch inside it sees
#          the GPU, then start an interactive shell in it with the course directory mounted
#          at /workspace/course and port 8888 published for JupyterLab
# Platform: spark
# Minimum memory: 8 GB
# Assumes: DGX OS with Docker and the NVIDIA Container Toolkit (preinstalled and configured on
#          DGX Spark per NVIDIA's documentation); your user in the docker group, or run with
#          sudo; about 12 GB of free disk for the image; run from the course directory
#          (~/llm-course)
#
# Usage: TAG=<yy.mm-py3> bash setup-env-spark.sh
#   TAG is a release tag from the container's NGC catalog page (26.08-py3 was the newest on
#   2026-09-12). The image is multi-arch, so the same tag serves the Spark's aarch64 CPU.
#   PyTorch and JupyterLab are in the image (NGC catalog page); the 26.08 release notes list
#   neither torchvision nor matplotlib, so the GPU check below also reports whether each
#   is present.
set -euo pipefail

TAG="${TAG:?set TAG to a release tag from the NGC catalog page, of the form yy.mm-py3}"
IMAGE="nvcr.io/nvidia/pytorch:${TAG}"
PORT="${PORT:-8888}"

command -v docker >/dev/null || { echo "setup-env-spark: docker is not installed" >&2; exit 1; }
if ! command -v nvidia-smi >/dev/null; then
  echo "setup-env-spark: nvidia-smi is not on PATH on the host; the GPU check below will fail" >&2
fi

echo "==> Pulling $IMAGE (about 11 GB compressed; the first pull takes a while)"
docker pull "$IMAGE"

echo "==> Checking that the GPU is visible to PyTorch inside the container"
docker run --gpus all --rm --interactive "$IMAGE" python - <<'PY'
import importlib.util
import torch
print("torch", torch.__version__, "cuda build:", torch.version.cuda)
if torch.cuda.is_available():
    print("accelerator: cuda /", torch.cuda.get_device_name(0))
else:
    print("accelerator: none visible")
for m in ("torchvision", "matplotlib"):
    print(f"{m}: " + ("present" if importlib.util.find_spec(m) else "missing; run: pip install " + m))
PY

echo "==> Starting an interactive shell with $PWD mounted at /workspace/course"
echo "    Port $PORT is published on 127.0.0.1 only; inside the container JupyterLab starts with:"
echo "      jupyter lab --ip 0.0.0.0 --port $PORT --no-browser --allow-root"
exec docker run --gpus all --interactive --tty --rm --ipc=host \
  --publish "127.0.0.1:${PORT}:${PORT}" \
  --volume "$PWD:/workspace/course" --workdir /workspace/course \
  "$IMAGE" bash