Support for Pyxis

New feature

Add native support for SLURM Pyxis container execution for GPU-enabled workloads

Use case

Nextflow is commonly used on HPC clusters with SLURM job schedulers. Many of these clusters, especially those focused on AI/ML workloads, use NVIDIA's Pyxis extension to run GPU-enabled containers directly through SLURM. Currently, Nextflow has excellent support for Singularity and Docker, but lacks native integration with Pyxis, requiring users to implement manual workarounds that bypass Nextflow's container abstraction.

This feature would benefit researchers and data scientists running GPU-intensive workflows on SLURM clusters where Pyxis is the preferred or only available container runtime. It would enable seamless container execution for deep learning, computational chemistry, and other GPU-accelerated pipelines.

Suggested implementation

1. Add a pyxis configuration block in nextflow.config similar to existing container engines:

   pyxis {
       enabled = true
       autoMounts = true
       // Additional Pyxis-specific options
   }
   

2. Extend the SlurmExecutor class to check for Pyxis configuration and modify how container directives are processed:

   • When a container directive is specified in a process and Pyxis is enabled, translate this to appropriate --container-image flags
   • Add container mount handling to ensure work directories are properly accessible

3. Implement automatic translation of Nextflow's container-related options to Pyxis-specific parameters:

   • Map volume mounts to --container-mounts
   • Support container working directory with --container-workdir
   • Handle other common container options

4. Ensure compatibility with SLURM array jobs and other advanced SLURM features when using Pyxis containers

5. Add documentation for the new Pyxis integration, with examples showing how to use GPU containers in SLURM environments

Note: An important implementation detail is that Pyxis specifically requires the use of srun to execute containers, unlike other container technologies that can work directly with sbatch. The implementation should handle this distinction automatically so users don't need to manually include srun commands in their process scripts.

Example of current workaround and issues

Here's a minimal example of what we currently need to do to use Pyxis, which is cumbersome and error-prone:

process runInPyxisContainer {
    cpus 2
    memory '4 GB'
    
    input:
    path inputFile
    
    output:
    path 'output.txt'
    
    script:
    """
    # Need to explicitly capture current directory
    CURRENT_DIR=\$(pwd)
    
    # Manually inject srun with container parameters
    srun --container-image=pytorch:24.02-py3 --container-mounts=\$CURRENT_DIR:\$CURRENT_DIR --container-workdir=\$CURRENT_DIR bash -c "
        python process_data.py $inputFile output.txt
    "
    """
}

Key issues I encountered with this approach:

1. File paths in symlinks are not properly handled when passed between processes
2. Need manual management of mount points for every directory needed
3. Hard to debug container startup issues
4. No integration with Nextflow's container directive
5. Process scripts become much more complex with this boilerplate
6. When using clusterOptions to pass --container-image, it gets sent to sbatch not srun, causing failures

A native implementation would let users simply specify container = 'pytorch:24.02-py3' and let Nextflow handle the Pyxis integration automatically.

Start by reading the SlurmExecutor and the existing container-engine handling, then review the proposed pyxis block in nextflow.config. Done means the container directive can translate to Pyxis options, mounts and working directories work, srun is handled automatically, SLURM array jobs remain compatible, and documentation includes GPU-container examples.