fix(slurm): make bootstrap container-safe

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

@@ -22,10 +22,18 @@ 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
+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` derives the cluster user path, installs the pinned Python
+under Lustre, 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,94 @@
+#!/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="${NSS_LUSTRE_DIR:-/lustre/fsw/portfolios/nemotron/projects/nemotron_data_dev/users/${user_name}}"
+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 user directory 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
 
@@ -34,52 +35,70 @@ 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
+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.
+The task derives the Slurm username from `id -un`, uses mise's pinned `uv`,
+installs the pinned Python under the user's 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.
 
-#### Nice to have
+Override the inferred user or Lustre directory only when the cluster account
+layout requires it:
+
+```bash
+NSS_SLURM_USER="your_user" MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
+NSS_LUSTRE_DIR="/custom/lustre/path" MISE_LOCKED=1 mise run bootstrap-nss-slurm cu129
+```
+
+Verify the postcondition before submitting jobs:
 
-- 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>
+readlink -f .venv/bin/python
+.venv/bin/safe-synthesizer --help >/dev/null
+```
+
+The Python path must start with `/lustre/fsw/` or its canonical `/lustre/fs11/`
+equivalent. `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 +274,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

@@ -104,15 +104,18 @@ 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.
+    if [[ ! -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
+    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"
@@ -129,7 +132,7 @@ else
     # 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] 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"

script/slurm/slurm_srun.sh

@@ -15,6 +15,18 @@ CPUS_PER_TASK="${CPUS_PER_TASK:-16}"
 SRUN_EXTRA="${SRUN_EXTRA:-}"
 export NEMO_DEPLOYMENT_TYPE="${NEMO_DEPLOYMENT_TYPE:-slurm-nvidia-internal}"
 
+# Pyxis gives variables baked into the container image precedence over the
+# Slurm step environment. Preserve the submitted job configuration explicitly
+# so a cached image created by another user cannot redirect paths or outputs.
+# Some Pyxis versions treat this as the complete container allowlist, so keep it
+# synchronized with job-critical variables consumed by the Slurm scripts.
+CONTAINER_ENV_VARS="${CONTAINER_ENV_VARS:-\
+USER_NAME,LUSTRE_DIR,NSS_SHARED_DIR,NSS_DIR,NSS_SLURM_DIR,CONFIG_DIR,\
+BASE_LOG_DIR,ADAPTER_PATH,NSS_ARTIFACTS_PATH,VLLM_CACHE_ROOT,HF_HOME,\
+UV_CACHE_DIR,UV_PYTHON_INSTALL_DIR,UV_PYTHON_BIN_DIR,UV_TOOL_DIR,\
+UV_CONCURRENT_DOWNLOADS,NSS_PYTHON_VERSION,WANDB_MODE,WANDB_PROJECT,\
+NEMO_DEPLOYMENT_TYPE,EXP_NAME,NSS_VERSION,PACKED_DATASETS,PACKED_CONFIGS,NSS_PHASE}"
+
 # No --time flag used in srun since we have a single srun in each allocation so
 # the time limit is controleld by the --time flag on the sbatch calls in
 # submit_slurm_jobs.sh.
@@ -23,6 +35,7 @@ srun -A "${ACCOUNT}" \
   --cpus-per-task="${CPUS_PER_TASK}" \
   --container-image="${CONTAINER_IMAGE}" \
   --container-mounts="${CONTAINER_MOUNTS}" \
+  --container-env="${CONTAINER_ENV_VARS}" \
   --export=ALL \
   ${SRUN_EXTRA} \
   /bin/bash -c "

script/slurm/submit_slurm_jobs.sh

@@ -154,13 +154,27 @@ prebuild_repo_venv() {
     flock -x 200
     unset VIRTUAL_ENV
     cd "${NSS_DIR}"
-    uv sync --frozen --extra cu129 --extra engine --group dev --python "${NSS_PYTHON_VERSION}"
+    MISE_IGNORED_CONFIG_PATHS="${MISE_IGNORED_CONFIG_PATHS:-${HOME}/.config/mise/config.toml}" \
+      MISE_LOCKED=1 \
+      NSS_SLURM_USER="${USER_NAME}" \
+      NSS_LUSTRE_DIR="${LUSTRE_DIR}" \
+      mise run bootstrap-nss-slurm cu129
   ) 200>"${lock_file}"
 
+  local lustre_root
+  local repo_python
+  lustre_root="$(readlink -f "${LUSTRE_DIR}")"
+  repo_python="$(readlink -f "${NSS_DIR}/.venv/bin/python" 2>/dev/null || true)"
+  if [[ -z "${repo_python}" || "${repo_python}" != "${lustre_root}/"* ]]; then
+    echo "[submit] ERROR: repo venv Python is not container-visible under Lustre: ${repo_python:-missing}" >&2
+    echo "[submit] Rebuild it: (cd ${NSS_DIR} && mise run bootstrap-nss-slurm cu129)" >&2
+    exit 1
+  fi
+
   # Functionally verify the venv by running safe-synthesizer once
   if ! "${NSS_DIR}/.venv/bin/safe-synthesizer" --help >/dev/null 2>&1; then
     echo "[submit] ERROR: repo venv at ${NSS_DIR}/.venv is not functional (safe-synthesizer failed to run)" >&2
-    echo "[submit] Rebuild it: (cd ${NSS_DIR} && rm -rf .venv && mise run bootstrap-nss cu129)" >&2
+    echo "[submit] Rebuild it: (cd ${NSS_DIR} && mise run bootstrap-nss-slurm cu129)" >&2
     exit 1
   fi
   echo "[submit] repo venv ready: $("${NSS_DIR}/.venv/bin/python" --version 2>&1)"