code snippets: comments as annotations, venv: reference to mksquashfs

Author: boeschf

docs/guides/storage.md

@@ -14,6 +14,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
 
 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.

docs/software/ml/pytorch.md

@@ -255,27 +255,28 @@ There are two ways to access the software provided by the uenv, once it has been
 
     !!! example "test mpi compilers and python provided by pytorch/v2.6.0"
         ```console
-        # start using the default view
-        $ uenv start pytorch/v2.6.0:v1 --view=default
+        $ uenv start pytorch/v2.6.0:v1 --view=default # (1)!
 
-        # the python executable provided by the uenv is the default, and is a recent version
-        $ which python
+        $ which python # (2)!
         /user-environment/env/default/bin/python
         $ python --version
         Python 3.13.0
 
-        # the mpi compiler wrappers are also available
-        $ which mpicc
+        $ which mpicc # (3)!
         /user-environment/env/default/bin/mpicc
         $ mpicc --version
         gcc (Spack GCC) 13.3.0
         $ gcc --version # the compiler wrapper uses the gcc provided by the uenv
         gcc (Spack GCC) 13.3.0
 
-        # exit the uenv
-        exit
+        $ exit # (4)!
         ```
 
+        1. start using the default view
+        2. the python executable provided by the uenv is the default, and is a recent version
+        3. the mpi compiler wrappers are also available
+        4. exit the uenv
+
 === "Spack"
 
     The pytorch uenv can also be used as a base for building software with
@@ -291,25 +292,36 @@ Uenvs are read-only, and cannot be modified. However, it is possible to add Pyth
 
 !!! example "creating a virtual environment on top of the uenv"
     ```console
-    # start the uenv
-    $ uenv start pytorch/v2.6.0:v1 --view=default
+    $ uenv start pytorch/v2.6.0:v1 --view=default # (1)!
 
-    # create a virtual environment
-    $ python -m venv ./my-venv
+    $ python -m venv ./my-venv # (2)!
 
-    # activate the virtual environment
-    $ source ./my-venv/bin/activate
+    $ source ./my-venv/bin/activate # (3)!
 
-    # install packages using pip
-    (my-venv) $ pip install <package>
+    (my-venv) $ pip install <package> # (4)!
 
-    # deactivate the virtual environment
-    (my-venv) $ deactivate
+    (my-venv) $ deactivate # (5)!
 
-    # exit the uenv
-    exit
+    $ 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
+    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.
+
 Alternatively one can use the uenv as [upstream Spack
 instance][ref-building-uenv-spack] to to add both Python and non-Python
 packages. However, this workflow is more involved and intended for advanced
@@ -326,36 +338,36 @@ Spack users.
     #SBATCH --ntasks-per-node=4
     #SBATCH --cpus-per-task=72
     #SBATCH --time=00:30:00
-    #SBATCH --uenv=pytorch/v2.6.0:/user-environment
+    #SBATCH --uenv=pytorch/v2.6.0:/user-environment # (1)!
     #SBATCH --view=default
 
     #################################
     # OpenMP environment variables #
     #################################
-    export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK  # (1)!
+    export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK  # (2)!
 
     #################################
     # PyTorch environment variables #
     #################################
-    export MASTER_ADDR=$(hostname) # (2)!
+    export MASTER_ADDR=$(hostname) # (3)!
     export MASTER_PORT=6000
     export WORLD_SIZE=$SLURM_NPROCS
-    export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 # (3)!
+    export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 # (4)!
 
     #################################
     # MPICH environment variables   #
     #################################
-    export MPICH_GPU_SUPPORT_ENABLED=0 # (4)!
+    export MPICH_GPU_SUPPORT_ENABLED=0 # (5)!
 
     #################################
     # CUDA environment variables    #
     #################################
-    export CUDA_CACHE_DISABLE=1 # (5)!
+    export CUDA_CACHE_DISABLE=1 # (6)!
 
     ############################################
     # NCCL and Fabric environment variables    #
     ############################################
-    export NCCL_NET="AWS Libfabric" # (6)!
+    export NCCL_NET="AWS Libfabric" # (7)!
     export NCCL_NET_GDR_LEVEL=PHB
     export NCCL_CROSS_NIC=1
     export FI_CXI_DISABLE_HOST_REGISTER=1
@@ -364,8 +376,8 @@ Spack users.
     export FI_CXI_DEFAULT_TX_SIZE=32768
     export FI_CXI_RX_MATCH_MODE=software
 
-    # (7)!
     # (8)!
+    # (9)!
     srun bash -c "
         export RANK=\$SLURM_PROCID
         export LOCAL_RANK=\$SLURM_LOCALID
@@ -374,20 +386,24 @@ Spack users.
     "
     ```
 
-    1. Only set `OMP_NUM_THREADS` if you are using OpenMP in your code.
-    2. These variables are used by PyTorch to initialize the distributed
+    1. The `--uenv` option is used to specify the uenv to use for the job. The
+       `--view=default` option is used to load all the packages provided by the
+       uenv.
+    2. Only set `OMP_NUM_THREADS` if you are using OpenMP in your code.
+    3. These variables are used by PyTorch to initialize the distributed
        backend. The `MASTER_ADDR` and `MASTER_PORT` variables are used to
-       determine the address and port of the master node. Additionally we also need
-       `RANK` and `LOCAL_RANK` but these must be set per-process, see below.
-    3. Enable more graceful exception handling, see [PyTorch
+       determine the address and port of the master node. Additionally we also
+       need `RANK` and `LOCAL_RANK` but these must be set per-process, see
+       below.
+    4. Enable more graceful exception handling, see [PyTorch
        documentation](https://pytorch.org/docs/stable/torch_nccl_environment_variables.html)
-    4. Disable GPU support in MPICH, as it [can lead to
+    5. Disable GPU support in MPICH, as it [can lead to
        deadlocks](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/mpi.html#inter-gpu-communication-with-cuda-aware-mpi)
        when using together with nccl.
-    5. Avoid writing JITed binaries to the (distributed) file system, which
+    6. Avoid writing JITed binaries to the (distributed) file system, which
        could lead to performance issues.
-    6. These variables should always be set for correctness and optimal
+    7. These variables should always be set for correctness and optimal
        performance when using NCCL, see [the detailed
        explanation][ref-communication-nccl].
-    7. `RANK` and `LOCAL_RANK` are set per-process by the SLURM job launcher.
-    8. Activate the virtual environment created on top of the uenv (if any).
+    8. `RANK` and `LOCAL_RANK` are set per-process by the SLURM job launcher.
+    9. Activate the virtual environment created on top of the uenv (if any).