dedicated uenv-venv section

Author: boeschf

docs/guides/storage.md

@@ -173,7 +173,7 @@ At first it can seem strange that a "high-performance" file system is significan
 Meta data lookups on Lustre are expensive compared to your laptop, where the local file system is able to aggressively cache meta data.
 
 [](){#ref-guides-storage-venv}
-### Python virtual environments with uenv
+### Squash Python virtual environments with uenv
 
 Python virtual environments can be very slow on Lustre, for example a simple `import numpy` command run on Lustre might take seconds, compared to milliseconds on your laptop.
 
@@ -191,7 +191,7 @@ This file can be mounted as a read-only [Squashfs](https://en.wikipedia.org/wiki
 
 #### Step 1: create the virtual environment
 
-The first step is to create the virtual environment using the usual workflow.
+The first step is to create the virtual environment using the usual workflow described in the [uenv documentation][ref-uenv-venv].
 
 === "uv"
 
@@ -205,13 +205,11 @@ The first step is to create the virtual environment using the usual workflow.
 
     # unset PYTHONPATH and set PYTHONUSERBASE to avoid conflicts
     unset PYTHONPATH
-    export PYTHONUSERBASE=/user-environment/env/default
+    export PYTHONUSERBASE="$(dirname "$(dirname "$(which python)")")"
 
     # create and activate a new relocatable venv using uv
     # in this case we explicitly select the python interpreter from the uenv view
-    uv venv -p /user-environment/env/default/bin/python --system-site-packages --relocatable --link-mode=copy /dev/shm/sqfs-demo/.venv
-    # You can also point to the uenv python with `uv venv -p $(which python) ...`
-    # which, among other things, enables user portability of the venv
+    uv venv --python $(which python) --system-site-packages --seed --relocatable --link-mode=copy /dev/shm/sqfs-demo/.venv
     cd /dev/shm/sqfs-demo
     source .venv/bin/activate
 

docs/software/ml/pytorch.md

@@ -779,55 +779,11 @@ There are two ways to access the software provided by the uenv, once it has been
 [](){#ref-uenv-pytorch-venv}
 ### Adding Python packages on top of the uenv
 
-Uenvs are read-only, and cannot be modified. However, it is possible to add Python packages on top of the uenv using virtual environments analogous to the setup with containers.
-
-```console title="Creating a virtual environment on top of the uenv"
-$ uenv start pytorch/v2.8.0:v1 --view=default # (1)!
-
-$ python -m venv --system-site-packages venv-uenv-pt2.8-v1 # (2)!
-
-$ source venv-uenv-pt2.8-v1/bin/activate # (3)!
-
-(venv-uenv-pt2.8-v1) $ pip install <package> # (4)!
-
-(venv-uenv-pt2.8-v1) $ deactivate # (5)!
-
-$ exit # (6)!
-```
-
-1. The `default` view is recommended, as it loads all the packages provided by the uenv.
-   This is important for PyTorch to work correctly, as it relies on the CUDA and NCCL libraries provided by the uenv.
-2. The virtual environment is created in the current working directory, and can be activated and deactivated like any other Python virtual environment.
-3. Activating the virtual environment will override the Python executable provided by the uenv, and use the one from the virtual environment instead.
-   This is important to ensure that the packages installed in the virtual environment are used.
-4. The virtual environment can be used to install any Python package.
-5. The virtual environment can be deactivated using the `deactivate` command.
-   This will restore the original Python executable provided by the uenv.
-6. The uenv can be exited using the `exit` command or by typing `ctrl-d`.
-
-
-!!! note "Squashing the virtual environment"
-    Python virtual environments can be slow on the parallel Lustre file system due to the amount of small files and potentially many processes accessing it.
-    If this becomes a bottleneck, consider [squashing the venv][ref-guides-storage-venv] into its own memory-mapped, read-only file system to enhance scalability and reduce load times.
-
-??? bug "Python packages from uenv view shadowing those in a virtual environment"
-    Some uenv views set the `PYTHONPATH` environment variable and/or do not set the `PYTHONUSERBASE` environment variable.
-    This can lead to unexpected behavior when using Python virtual environments on top of the uenv, as the packages installed in the uenv view may take precedence over those in the virtual environment.
-    A possible workaround is to unset the `PYTHONPATH` and set the `PYTHONUSERBASE` environment variables, as described in the [Python virtual environments with uenv guide][ref-guides-storage-venv]:
-    ```bash
-    unset PYTHONPATH
-    export PYTHONUSERBASE=/user-environment/env/default
-    ```
-    It is recommended to apply this workaround if you are constrained by a Python package version installed in the uenv view that you need to change for your application.
-
-!!! note
-    Keep in mind that
-
-     * this virtual environment is _specific_ to this particular uenv and won't actually work unless you are using it from inside this uenv - it relies on the resources packaged inside the uenv.
-     * every Slurm job making use of this virtual environment will need to activate it first (_inside_ the `srun` command). 
-
-Alternatively one can use the uenv as [upstream Spack instance][ref-build-uenv-spack] to to add both Python and non-Python packages.
-However, this workflow is more involved and intended for advanced Spack users.
+Python virtual environments can be created on top of the uenv to install additional Python packages not provided by the uenv itself, or to override existing packages.
+Please refer to the [Python virtual environments with uenv documentation][ref-uenv-venv] and the [guide on performance][ref-guides-storage-venv] for more details on
+- creating and managing virtual environments on top of uenvs
+- best practices and caveats when using virtual environments with uenvs
+- troubleshooting common issues
 
 ### Running PyTorch jobs with Slurm
 

docs/software/uenv/index.md

@@ -775,5 +775,108 @@ $ which uenv
 /usr/bin/uenv
 ```
 
+[](){#ref-uenv-venv}
+## Python virtual environments on top of uenv views
 
+When stacking a Python virtual environment on top of a _uenv view_, keep Python’s import resolution predictable with the following:
 
+- **Unset `PYTHONPATH`**. Anything there is *prepended* to Python's `sys.path`, which can lead to surprising imports.
+- **Set `PYTHONUSERBASE` to the view's root directory** (e.g., `/user-environment/env/default`) so the interpreter’s _user site_ resolves inside the view.
+  - You can derive this automatically from the interpreter you’re about to use: take the parent of `which python`:
+    ```bash
+    export PYTHONUSERBASE=$(dirname $(dirname $(which python)))
+    ```
+  - Do not use tools that resolve symlinks (such as `readlink -f` or Python's `Path.resolve()`), as the Python interpreter in the _uenv view_ is a symlink - following it would point outside the view.
+- **Create the venv with `--system-site-packages`**.
+`venv` disables the user site by default; enabling system site restores both the system site and the user site, so packages provided by the _uenv view_ become visible inside the venv.
+
+=== "uv"
+
+    ```console title="Creating a Python virtual environment on top of a uenv view"
+    # start the uenv with the default view
+    $ uenv start --view=default prgenv-gnu/25.6:v2
+
+    # unset PYTHONPATH to avoid surprises
+    $ unset PYTHONPATH
+
+    # set PYTHONUSERBASE to the root of the view
+    $ export PYTHONUSERBASE="$(dirname "$(dirname "$(which python)")")"
+
+    # create the virtual environment with access to system site packages
+    # - optionally seed it with pip, setuptools and wheel
+    # - optionally make it relocatable and copy linked files (useful for moving venvs)
+    $ uv venv --python $(which python) --system-site-packages --seed --relocatable --link-mode=copy path/to/my-venv
+
+    # activate the virtual environment
+    $ source path/to/my-venv/bin/activate
+
+    # verify that packages from the uenv are visible (note Locations)
+    (my-venv) $ python -m pip list -v
+    Package Version Location                                                   Installer
+    ------- ------- ---------------------------------------------------------- ---------
+    meson   1.7.0   /user-environment/env/default/lib/python3.13/site-packages pip
+    pip     25.3    /path/to/my-venv/lib/python3.13/site-packages              uv
+
+    # upgrade a package into the venv (overrides the view's version)
+    (my-venv) $ uv pip install --upgrade meson
+
+    # verify that the upgraded package is now coming from the venv
+    (my-venv) $ python -m pip list -v
+    Package Version Location                                         Installer
+    ------- ------- ------------------------------------------------ ---------
+    meson   1.9.1   /path/to/my-venv-uv/lib/python3.13/site-packages uv
+    pip     25.3    /path/to/my-venv-uv/lib/python3.13/site-packages uv
+    ```
+
+=== "venv"
+
+    ```console title="Creating a Python virtual environment on top of a uenv view"
+    # start the uenv with the default view
+    $ uenv start --view=default prgenv-gnu/25.6:v2
+
+    # unset PYTHONPATH to avoid surprises
+    $ unset PYTHONPATH
+
+    # set PYTHONUSERBASE to the root of the view
+    $ export PYTHONUSERBASE="$(dirname "$(dirname "$(which python)")")"
+
+    # create the virtual environment with access to system site packages
+    $ python -m venv --system-site-packages path/to/my-venv
+
+    # activate the virtual environment
+    $ source path/to/my-venv/bin/activate
+
+    # verify that packages from the uenv are visible (note Locations)
+    (my-venv) $ python -m pip list -v
+    python -m pip list -v
+    Package Version Location                                                   Installer
+    ------- ------- ---------------------------------------------------------- ---------
+    meson   1.7.0   /user-environment/env/default/lib/python3.13/site-packages pip
+    pip     25.1.1  /path/to/my-venv/lib/python3.13/site-packages              pip
+
+    # upgrade a package into the venv (overrides the view's version)
+    (my-venv) $ pip install --upgrade meson
+
+    # verify that the upgraded package is now coming from the venv
+    (my-venv) $ pip list
+    Package Version Location                                      Installer
+    ------- ------- --------------------------------------------- ---------
+    meson   1.9.1   /path/to/my-venv/lib/python3.13/site-packages pip
+    pip     25.1.1  /path/to/my-venv/lib/python3.13/site-packages pip
+    ```
+
+!!! note "Listing only the packages that live in the uenv’s user-site"
+    To see _just_ what's in the uenv view (not what's installed into the venv), list by the user-site path that your interpreter is using:
+    ```console
+    (my-venv) $ python -m pip list -v --path "$(python -c 'import site; print(site.getusersitepackages())')"
+    ```
+
+!!! note "Troubleshooting & Gotchas"
+    - `pip install --user` will fail here.
+    The uenv is a read-only squashfs; a `--user` install would try to write into `PYTHONUSERBASE` (the uenv), which is not possible.
+    - Some uenv views already set `PYTHONUSERBASE`. If you start a uenv view that does this, you can skip setting `PYTHONUSERBASE` yourself.
+    - The virtual environment is _specific_ to a particular uenv and won't work unless used from inside this excact uenv - it relies on the resources packaged inside the uenv.
+
+!!! note "Performance considerations"
+    On our Lustre parallel file system, large venvs can be slow due to many small files.
+    See [How to squash virtual environments][ref-guides-storage-venv] for turning a venv into a compact image to improve startup and import performance.