From 517a40bf0d3cf2459e2500aa17b7a0f4261ccc2d Mon Sep 17 00:00:00 2001
From: Daniele Nerini <daniele.nerini@meteoswiss.ch>
Date: Thu, 10 Sep 2026 11:58:07 +0200
Subject: [PATCH] Mount venv squashfs via srun --uenv instead of squashfs-mount
 (#260)

On Linard, the cluster now provides a native `srun
--uenv=<squashfs>:<mountpoint>` flag (via the `slurm-uenv-mount` SPANK
plugin), which mounts the squashfs inside the task's own namespace on
the allocated node.

---------

Co-authored-by: Luca Lanzilao <45151790+lclanzi@users.noreply.github.com>
---
 workflow/rules/inference.smk | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk
index 1a83b57c..c2f90ec6 100644
--- a/workflow/rules/inference.smk
+++ b/workflow/rules/inference.smk
@@ -293,7 +293,7 @@ rule inference_execute:
             cd {params.workdir}
 
             _run_inference() {{
-                local VENV=$1
+                local VENV=/user-environment
                 source "$VENV/bin/activate"
 
                 if [ "{params.disable_local_definitions}" = "False" ]; then
@@ -307,19 +307,21 @@ rule inference_execute:
                     CMD_ARGS+=(runner.parallel.cluster=slurm)
                 fi
 
-                srun \
-                    --unbuffered \
-                    --partition={resources.slurm_partition} \
-                    --cpus-per-task={resources.cpus_per_task} \
-                    --mem-per-cpu={resources.mem_mb_per_cpu} \
-                    --time={resources.runtime} \
-                    --gres={resources.gres} \
-                    --ntasks={resources.ntasks} \
-                    anemoi-inference run config.yaml "${{CMD_ARGS[@]}}"
+                anemoi-inference run config.yaml "${{CMD_ARGS[@]}}"
             }}
             export -f _run_inference
 
-            squashfs-mount {params.env_path}:/user-environment -- bash -c '_run_inference /user-environment'
+            srun \
+                --job-name=anemoi-inference \
+                --uenv={params.env_path}:/user-environment \
+                --unbuffered \
+                --partition={resources.slurm_partition} \
+                --cpus-per-task={resources.cpus_per_task} \
+                --mem-per-cpu={resources.mem_mb_per_cpu} \
+                --time={resources.runtime} \
+                --gres={resources.gres} \
+                --ntasks={resources.ntasks} \
+                bash -c '_run_inference'
         ) >{log} 2>&1
         touch {output.okfile}
         """