Allow unprivileged docker runs for local testing (#380)

Author: p-zubieta

.github/workflows/docker-ci.yml

@@ -29,6 +29,8 @@ jobs:
         with:
           context: .
           tags: pysages:latest
+          build-args: |
+            JAX_PLATFORMS=cpu
           outputs: type=docker,dest=/tmp/pysages.tar
       - name: Upload artifact
         uses: actions/upload-artifact@v4
@@ -142,11 +144,11 @@ jobs:
       - name: Load and run test
         run: |
           SCRIPT="cd PySAGES/examples/openmm/umbrella_integration &&"
-          SCRIPT="${SCRIPT} mpirun --allow-run-as-root --oversubscribe -n 6"
+          SCRIPT="${SCRIPT} mpirun --oversubscribe -n 6"
           SCRIPT="${SCRIPT} python3 -m mpi4py.futures"
           SCRIPT="${SCRIPT} integration.py --replicas=5 --time-steps=1000 --mpi"
           docker load --input /tmp/pysages.tar
-          docker run --privileged -t pysages bash -c "${SCRIPT}"
+          docker run --user pysages -e HOME=/home/pysages -t pysages bash -c "${SCRIPT}"
 
   cv-pucker-openmm-unbiased:
     runs-on: ubuntu-latest
@@ -214,11 +216,11 @@ jobs:
       - name: Load and run test
         run: |
           SCRIPT="cd PySAGES/examples/openmm/string &&"
-          SCRIPT="${SCRIPT} mpirun --allow-run-as-root --oversubscribe -n 6"
+          SCRIPT="${SCRIPT} mpirun --oversubscribe -n 6"
           SCRIPT="${SCRIPT} python3 -m mpi4py.futures"
           SCRIPT="${SCRIPT} integration.py --replicas=5 --string-steps=2 --time-steps=1000 --mpi"
           docker load --input /tmp/pysages.tar
-          docker run --privileged -t pysages bash -c "${SCRIPT}"
+          docker run --user pysages -e HOME=/home/pysages -t pysages bash -c "${SCRIPT}"
 
   openmm-nacl-metad-displacement:
     runs-on: ubuntu-latest
@@ -396,12 +398,12 @@ jobs:
         run: |
           SCRIPT="cd PySAGES/examples/hoomd-blue/umbrella_integration &&"
           SCRIPT="${SCRIPT} python3 ./gen_gsd.py &&"
-          SCRIPT="${SCRIPT} mpirun --allow-run-as-root --oversubscribe -n 2"
+          SCRIPT="${SCRIPT} mpirun --oversubscribe -n 2"
           SCRIPT="${SCRIPT} python3 -m mpi4py.futures"
           SCRIPT="${SCRIPT} integration.py --replicas=5 --time-steps=1000 --mpi &&"
           SCRIPT="${SCRIPT} mkdir /tmp/plots && mv *.pdf /tmp/plots/"
           docker load --input /tmp/pysages.tar
-          docker run -v /tmp:/tmp -t pysages bash -c "${SCRIPT}"
+          docker run -v /tmp:/tmp --user pysages -e HOME=/home/pysages -t pysages bash -c "${SCRIPT}"
       - name: Upload artifacts
         uses: actions/upload-artifact@v4
         with:

Dockerfile

@@ -6,8 +6,17 @@ RUN python -m pip install --upgrade pip
 RUN python -m pip install ase dill "gsd<3.3" matplotlib "pyparsing<3"
 
 # Install JAX and JAX-based libraries
-RUN python -m pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-RUN python -m pip install --upgrade "dm-haiku<0.0.11" "e3nn-jax!=0.20.4" "jax-md>=0.2.7" jaxopt
+# TODO: Remove jax pin once DLPack buffer alignment is fixed
+RUN python -m pip install --upgrade "jax[cuda]==0.4.34" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN python -m pip install --upgrade "jax==0.4.34" "dm-haiku<0.0.11" "e3nn-jax!=0.20.4" "jax-md>=0.2.7" jaxopt
+
+ARG JAX_PLATFORMS=""
+ENV JAX_PLATFORMS=${JAX_PLATFORMS}
 
 COPY . /PySAGES
 RUN pip install /PySAGES/
+
+# Disable CAP_SYS_PTRACE requirement for vader's CMA single-copy path.
+ENV OMPI_MCA_btl_vader_single_copy_mechanism=none
+# Non-privileged setups (e.g. MPI CI jobs) can use `docker run --user pysages`.
+RUN useradd -m -u 1000 pysages && chown -R pysages:pysages /PySAGES