fix(slurm): make bootstrap container-safe

.agents/skills/uv-build/SKILL.md

@@ -22,10 +22,19 @@ mise run bootstrap-nss cpu       # + engine + CPU PyTorch
 mise run bootstrap-nss cu129     # + engine + CUDA 12.9 PyTorch
 mise run bootstrap-nss cuda      # alias for cu129
 mise run bootstrap-nss engine    # + engine (no torch)
+
+# Slurm: force Python, caches, and the project venv onto Lustre
+LUSTRE_DIR="/path/to/container-visible/project/directory" \
+  MISE_IGNORED_CONFIG_PATHS="$HOME/.config/mise/config.toml" \
+  MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
 ```
 
 Under the hood: `uv sync --frozen --extra <extra> [--extra engine] --group dev`
 
+`bootstrap-nss-slurm` requires `LUSTRE_DIR`, installs the pinned Python under
+that directory, recreates `.venv` if its interpreter is not container-visible,
+then runs the same frozen profile sync as `bootstrap-nss`.
+
 ## Extras and Conflicts
 
 | Extra | What it installs |

.mise/tasks/_lib.sh

@@ -38,6 +38,30 @@ resolve_venv_path() {
   printf '%s\n' "${venv_path}"
 }
 
+sync_nss_dependencies() {
+  local extra="${1:?Python dependency profile is required}"
+  local sync_uv_bin="${NSS_UV_BIN:-uv}"
+
+  case "${extra}" in
+    cuda|cu129)
+      "${sync_uv_bin}" sync --frozen --extra cu129 --extra engine --group dev
+      ;;
+    cpu)
+      "${sync_uv_bin}" sync --frozen --extra cpu --extra engine --group dev
+      ;;
+    engine)
+      "${sync_uv_bin}" sync --frozen --extra engine --group dev
+      ;;
+    dev)
+      "${sync_uv_bin}" sync --frozen --group dev
+      ;;
+    *)
+      echo "Error: Invalid extra '${extra}'. Use one of: dev engine cpu cuda cu129" >&2
+      return 1
+      ;;
+  esac
+}
+
 resolve_container_cmd() {
   local container_cmd="${CONTAINER_CMD:-$(command -v podman 2>/dev/null || command -v docker 2>/dev/null || true)}"
 

.mise/tasks/bootstrap-nss

@@ -7,27 +7,10 @@
 #USAGE }
 
 set -euo pipefail
+source "${MISE_CONFIG_ROOT}/.mise/tasks/_lib.sh"
 
 mise run venv
 
 echo "~~~~~~"
 echo "attempting to install nss package with primary extra: ${usage_extra?}"
-
-case "${usage_extra?}" in
-  cuda|cu129)
-    uv sync --frozen --extra cu129 --extra engine --group dev
-    ;;
-  cpu)
-    uv sync --frozen --extra cpu --extra engine --group dev
-    ;;
-  engine)
-    uv sync --frozen --extra engine --group dev
-    ;;
-  dev)
-    uv sync --frozen --group dev
-    ;;
-  *)
-    echo "Error: Invalid extra '${usage_extra?}'. Use one of: dev engine cpu cuda cu129" >&2
-    exit 1
-    ;;
-esac
+sync_nss_dependencies "${usage_extra?}"

.mise/tasks/bootstrap-nss-slurm

@@ -0,0 +1,99 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#MISE description="Install project dependencies with Python, caches, and the virtualenv on Lustre for Slurm containers."
+#USAGE arg "<extra>" help="Python dependency profile" {
+#USAGE   choices "dev" "engine" "cpu" "cuda" "cu129"
+#USAGE }
+
+set -euo pipefail
+source "${MISE_CONFIG_ROOT}/.mise/tasks/_lib.sh"
+
+user_name="${NSS_SLURM_USER:-${USER_NAME:-$(id -un)}}"
+lustre_dir="${LUSTRE_DIR:-}"
+if [[ -z "${lustre_dir}" ]]; then
+  echo "Error: LUSTRE_DIR must be set to the user's container-visible project directory" >&2
+  exit 1
+fi
+
+python_version="${NSS_PYTHON_VERSION:-$(resolve_python_version)}"
+venv_path="$(resolve_venv_path)"
+uv_bin="$(mise which uv)"
+readonly user_name lustre_dir python_version venv_path uv_bin
+
+if [[ ! -d "${lustre_dir}" ]]; then
+  echo "Error: LUSTRE_DIR does not exist: ${lustre_dir}" >&2
+  exit 1
+fi
+
+lustre_real="$(readlink -f "${lustre_dir}")"
+venv_parent_real="$(readlink -f "$(dirname "${venv_path}")")"
+readonly lustre_real venv_parent_real
+if [[ "${venv_parent_real}" != "${lustre_real}" && "${venv_parent_real}" != "${lustre_real}/"* ]]; then
+  echo "Error: Slurm virtualenv must be located under Lustre: ${venv_path}" >&2
+  exit 1
+fi
+
+export USER_NAME="${user_name}"
+export LUSTRE_DIR="${lustre_dir}"
+export NSS_DIR="${MISE_CONFIG_ROOT}"
+export UV_CACHE_DIR="${LUSTRE_DIR}/.cache/uv"
+export UV_PYTHON_INSTALL_DIR="${LUSTRE_DIR}/.local/share/uv/python"
+export UV_PYTHON_BIN_DIR="${LUSTRE_DIR}/.local/bin"
+export UV_TOOL_DIR="${LUSTRE_DIR}/.local/share/uv/tools"
+export HF_HOME="${LUSTRE_DIR}/.cache/huggingface"
+export VLLM_CACHE_ROOT="${LUSTRE_DIR}/.cache/vllm"
+export NSS_UV_BIN="${uv_bin}"
+
+mkdir -p \
+  "${UV_CACHE_DIR}" \
+  "${UV_PYTHON_INSTALL_DIR}" \
+  "${UV_PYTHON_BIN_DIR}" \
+  "${UV_TOOL_DIR}" \
+  "${HF_HOME}" \
+  "${VLLM_CACHE_ROOT}"
+
+echo "[slurm-bootstrap] installing Python ${python_version} under ${UV_PYTHON_INSTALL_DIR}"
+"${uv_bin}" python install --no-bin --install-dir "${UV_PYTHON_INSTALL_DIR}" "${python_version}"
+
+python_bin="$(PATH=/usr/bin:/bin VIRTUAL_ENV= \
+  "${uv_bin}" python find --managed-python --no-project "${python_version}")"
+readonly python_bin
+if [[ ! -x "${python_bin}" ]]; then
+  echo "Error: uv did not resolve a managed Lustre Python executable: ${python_bin}" >&2
+  exit 1
+fi
+
+expected_python="$(readlink -f "${python_bin}")"
+readonly expected_python
+if [[ "${expected_python}" != "${lustre_real}/"* ]]; then
+  echo "Error: managed Python is not container-visible under Lustre: ${expected_python}" >&2
+  exit 1
+fi
+
+if [[ -e "${venv_path}" ]]; then
+  current_python="$(readlink -f "${venv_path}/bin/python" 2>/dev/null || true)"
+  if [[ "${current_python}" != "${expected_python}" ]]; then
+    echo "[slurm-bootstrap] recreating ${venv_path}; current Python is ${current_python:-unavailable}"
+    rm -rf "${venv_path}"
+  fi
+fi
+
+"${uv_bin}" venv --seed --allow-existing --python "${python_bin}" "${venv_path}"
+
+echo "[slurm-bootstrap] syncing NSS dependencies with profile ${usage_extra?}"
+sync_nss_dependencies "${usage_extra?}"
+
+venv_python="$(readlink -f "${venv_path}/bin/python")"
+readonly venv_python
+if [[ "${venv_python}" != "${lustre_real}/"* ]]; then
+  echo "Error: virtualenv Python is not container-visible under Lustre: ${venv_python}" >&2
+  exit 1
+fi
+
+if ! "${venv_path}/bin/safe-synthesizer" --help >/dev/null 2>&1; then
+  echo "Error: Safe Synthesizer entry point is not usable in ${venv_path}" >&2
+  exit 1
+fi
+
+echo "[slurm-bootstrap] ready: ${venv_python}"

script/slurm/README.md

@@ -13,12 +13,13 @@ Jobs are submitted via `submit_slurm_jobs.sh`, which launches a containerized `s
 - `submit_slurm_jobs.sh`: Submits Slurm array jobs for each config and dataset. Supports two-stage TRAIN→GEN pipeline.
 - `slurm_nss_matrix.sh`: Picks dataset and config and launches the python entrypoint inside the container. Honors `NSS_PHASE=train|generate|end_to_end`.
 - `slurm_srun.sh`: Wraps `srun` with container image and mounts, mostly just a pass through, primary logic is in `submit_slurm_jobs.sh` and `slurm_nss_matrix.sh`.
+- `.mise/tasks/bootstrap-nss-slurm`: Installs a container-visible Python and project virtualenv under the current user's Lustre directory.
 - `configs/*.yaml`: Major configs we support. Use the config basenames from this directory in commands (for example, `smollm3-nodp`, `smollm3-dp`, etc.). The current set is the cross product of 3 pre-trained models and 2 DP settings (on or off).
 
-Pipeline entrypoints (invoked by Slurm scripts) via uv:
-- `uv run safe-synthesizer run --run-path <path>` (full end-to-end pipeline)
-- `uv run safe-synthesizer run train --run-path <path>` (PII replacement + training only)
-- `uv run safe-synthesizer run generate --run-path <path>` (generation + evaluation only)
+Pipeline entrypoints invoked from the prebuilt project virtualenv:
+- `.venv/bin/safe-synthesizer run --run-path <path>` (full end-to-end pipeline)
+- `.venv/bin/safe-synthesizer run train --run-path <path>` (PII replacement + training only)
+- `.venv/bin/safe-synthesizer run generate --run-path <path>` (generation + evaluation only)
 
 ### Prerequisites
 
@@ -27,59 +28,88 @@ Pipeline entrypoints (invoked by Slurm scripts) via uv:
 - Weights & Biases API Key: W&B logging is enabled by default (`WANDB_MODE=online`). You will need a `WANDB_API_KEY` — request an account [here](https://confluence.nvidia.com/display/AIALGO/Weights+and+Biases+%28WandB%29+Enterprise+Account). Set `WANDB_MODE=disabled` in `env_variables.sh` to skip W&B.
 - Enroot Credentials: Follow https://confluence.nvidia.com/display/HWINFCSSUP/Using+Containers#UsingContainers-SettingupEnrootCredentials. You should add the lines for all 3 of `nvcr.io`, `authn.nvidia.com`, and `gitlab-master.nvidia.com`.
 - Clone Safe-Synthesizer
+
+The instructions below assume that Safe-Synthesizer is cloned directly under
+`LUSTRE_DIR` and that commands after cloning are run from the repository root.
+
 ```bash
 export USER_NAME="$USER" # Or hardcode username in slurm
 export LUSTRE_DIR="/lustre/fsw/portfolios/nemotron/projects/nemotron_data_dev/users/${USER_NAME}"
-cd $LUSTRE_DIR
+cd "${LUSTRE_DIR}"
 git clone git@github.com:NVIDIA-NeMo/Safe-Synthesizer.git
 cd Safe-Synthesizer
 ```
-- uv and python install in the slurm cluster
-  - DO NOT FOLLOW the general CONTRIBUTING.md or README.md instructions for installation and setup, unless you understand exactly what's being installed where and how that interacts with the distributed nature of a slurm cluster.
-  - The following setup is strongly recommended, but is not the only way to get things working.
-  - The key issues about working in slurm we need to address
-    - /home/$USER is quite small (10 GB) and not recommended for accessing data (easily filled up by uv cache)
-    - Slurm jobs may run in containers with different $HOME locations (and different users/uids)
-  - Thus we put uv and python in your user directory in /lustre and not in /home/$USER
+
+#### Bootstrap the Slurm Python environment
+
+Do not use the general development bootstrap for Slurm. Slurm containers mount
+`/lustre`, but may not have access to the login node's `/home` directory. The
+Python interpreter, project virtualenv, package caches, and model caches must
+therefore resolve under Lustre.
+
+From the repository root, install the repository-pinned mise tools and run the
+Slurm-specific bootstrap:
+
 ```bash
-export USER_NAME="$USER" # Or hardcode username in slurm
-export LUSTRE_DIR="/lustre/fsw/portfolios/nemotron/projects/nemotron_data_dev/users/${USER_NAME}"
-# Install `uv` to your lustre directory
-curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="$LUSTRE_DIR/.uv/bin" sh
-# Set environment variables so `uv` uses $LUSTRE_DIR subdirectories for storage
-export UV_CACHE_DIR="${LUSTRE_DIR}/.cache/uv"
-export UV_PYTHON_INSTALL_DIR="${LUSTRE_DIR}/.local/share/uv/python"
-export UV_PYTHON_BIN_DIR="${LUSTRE_DIR}/.local/bin"
-export UV_TOOL_DIR="${LUSTRE_DIR}/.local/share/uv/tools"
-# With the above env vars, the usual mise task should work.
-# Note this may be quite slow the first time due to very slow network
-# connectivity on slurm to download from pypi, but subsequent executions
-# (such as startup for your jobs) should be much faster since uv will
-# pull cached wheels from UV_CACHE_DIR.
-# (Be sure to run from the root of the Safe-Synthesizer repo)
-mise run bootstrap-nss cu129
+make install-mise
+export PATH="${HOME}/.local/bin:${PATH}"
+export MISE_IGNORED_CONFIG_PATHS="${HOME}/.config/mise/config.toml"
+mise install --locked
+MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
 ```
 
-`cu129` and `cuda` select the same CUDA 12.9 dependency profile.
+Existing Slurm checkouts should run this bootstrap once after pulling the
+change. Scheduled launch environments, including GitLab jobs, must also add
+`${HOME}/.local/bin` to `PATH` before invoking `submit_slurm_jobs.sh`.
 
-#### Nice to have
+The task requires the exported `LUSTRE_DIR`, derives the Slurm username from
+`id -un`, uses mise's pinned `uv`, installs the pinned Python under that Lustre
+directory, and creates `.venv` with that exact interpreter. If an existing
+`.venv` points outside Lustre, the task recreates it. Existing Python and uv
+installations under `/home` do not need to be removed.
+
+Override the inferred user or select a different Lustre project directory when
+the cluster account layout requires it:
 
-- Passwordless login See https://confluence.nvidia.com/display/HWINFCSSUP/Setting+Up+Passwordless+SSH+Key+Authentication?src=contextnavpagetreemode
-- Env vars in .bashrc
-  - Add `export VARIABLE=VALUE` to the end of `~/.bashrc` for commonly used environment variables, like `USER_NAME` and `LUSTRE_DIR`.
-  - Recommended snippet to have in `~/.bashrc` so uv and python work on login node and slurm jobs:
 ```bash
-export USER_NAME=<your slurm user name>
+NSS_SLURM_USER="your_user" MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
+LUSTRE_DIR="/custom/lustre/path" MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
+```
+
+Verify the postcondition before submitting jobs:
+
+```bash
+readlink -f .venv/bin/python
+.venv/bin/safe-synthesizer --help >/dev/null
+```
+
+The Python path must resolve beneath the canonical path reported by
+`readlink -f "${LUSTRE_DIR}"`. `cu129` and `cuda` select the same CUDA 12.9
+dependency profile.
+
+Repo mode, the default when `--nss-version` is omitted, does not require a
+separate uv installation under Lustre. PyPI mode uses uv inside the container
+and still requires `${LUSTRE_DIR}/.uv/bin/env`.
+
+For PyPI mode only, install the same uv version that mise has pinned:
+
+```bash
+export USER_NAME="${USER_NAME:-$(id -un)}"
 export LUSTRE_DIR="/lustre/fsw/portfolios/nemotron/projects/nemotron_data_dev/users/${USER_NAME}"
+uv_version="$(mise exec -- uv --version | awk '{print $2}')"
+curl -LsSf "https://astral.sh/uv/${uv_version}/install.sh" \
+  | env UV_INSTALL_DIR="${LUSTRE_DIR}/.uv/bin" sh
+```
+
+#### Nice to have
 
-# (May be added automatically by uv)
-. "${LUSTRE_DIR}/.uv/bin/env"
+- Passwordless login See https://confluence.nvidia.com/display/HWINFCSSUP/Setting+Up+Passwordless+SSH+Key+Authentication?src=contextnavpagetreemode
+- Env vars in `.bashrc`
+  - This is optional for bootstrap, but avoids exporting the submission user in every shell.
 
-export UV_CACHE_DIR="${LUSTRE_DIR}/.cache/uv"
-export UV_PYTHON_INSTALL_DIR="${LUSTRE_DIR}/.local/share/uv/python"
-export UV_PYTHON_BIN_DIR="${LUSTRE_DIR}/.local/bin"
-export UV_TOOL_DIR="${LUSTRE_DIR}/.local/share/uv/tools"
-export HF_HOME="${LUSTRE_DIR}/.cache/huggingface"
+```bash
+export USER_NAME="<your slurm user name>"
+export LUSTRE_DIR="/lustre/fsw/portfolios/nemotron/projects/nemotron_data_dev/users/${USER_NAME}"
 ```
 
 
@@ -255,6 +285,7 @@ W&B is enabled by default with `WANDB_MODE=online` in `env_variables.sh`. Make s
 - Missing token file/key: create `${LUSTRE_DIR}/.api_tokens.sh` with `NSS_INFERENCE_KEY` and `chmod 600`.
 - Missing config files: verify `CONFIGS` in `env_variables.sh` and files in `CONFIG_DIR`.
 - Permission errors: confirm your `/lustre/.../${USER_NAME}` paths and file perms.
+- Virtualenv Python under `/home`: run `MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129`; the task recreates `.venv` with a Lustre-managed interpreter.
 
 #### cpu bind errors
 

script/slurm/slurm_nss_matrix.sh

@@ -94,6 +94,14 @@ else
     config=${all_configs[0]}
 fi
 
+# Fail before installing system packages when PyPI mode cannot access uv.
+if [[ -n "${NSS_VERSION:-}" && ! -f "${LUSTRE_DIR}/.uv/bin/env" ]]; then
+    echo "[NSS SLURM] ERROR: PyPI mode requires uv at ${LUSTRE_DIR}/.uv/bin" >&2
+    echo "[NSS SLURM] Install uv as described in script/slurm/README.md," \
+        "or omit --nss-version to use repo mode." >&2
+    exit 1
+fi
+
 # Ensure minimal build toolchain inside container (no-op if already present)
 apt-get update && apt-get install -y --no-install-recommends \
         curl \
@@ -104,15 +112,13 @@ apt-get update && apt-get install -y --no-install-recommends \
         libcurl4 \
         libcurl3-gnutls
 
-# Ensure Python environment is available inside the container
-source "${LUSTRE_DIR}/.uv/bin/env"
-
 # NSS_PYTHON_VERSION is resolved in env_variables.sh (from .python-version, with
 # an optional override), which is always sourced before this script runs.
 
 if [[ -n "${NSS_VERSION:-}" ]]; then
     # Install nemo-safe-synthesizer from PyPI into a versioned venv cached on
     # lustre so concurrent array jobs can share it without redundant downloads.
+    source "${LUSTRE_DIR}/.uv/bin/env"
     PYPI_VENV="${LUSTRE_DIR}/.venv_nss_py${NSS_PYTHON_VERSION}_${NSS_VERSION}"
     uv venv --python "${NSS_PYTHON_VERSION}" "${PYPI_VENV}"
     source "${PYPI_VENV}/bin/activate"
@@ -128,8 +134,10 @@ else
     # submit_slurm_jobs.sh, rather than each array task running `uv sync` against
     # the shared Lustre .venv (which could race and corrupt it).
     if [[ ! -x "${NSS_DIR}/.venv/bin/safe-synthesizer" ]]; then
-        echo "[NSS SLURM] ERROR: no usable repo venv at ${NSS_DIR}/.venv (missing safe-synthesizer entry point)" >&2
-        echo "[NSS SLURM] Submit via submit_slurm_jobs.sh (it builds the venv once), or build it manually: (cd ${NSS_DIR} && mise run bootstrap-nss cu129)" >&2
+        echo "[NSS SLURM] ERROR: no usable repo venv at ${NSS_DIR}/.venv" \
+            "(missing safe-synthesizer entry point)" >&2
+        echo "[NSS SLURM] Submit via submit_slurm_jobs.sh (it builds the venv once)," \
+            "or build it manually: (cd ${NSS_DIR} && mise run bootstrap-nss-slurm cu129)" >&2
         exit 1
     fi
     echo "[NSS SLURM] Using pre-built repo venv at ${NSS_DIR}/.venv"