updated container files

docker/Containerfile-test

@@ -0,0 +1,35 @@
+# podman-hpc login docker.io
+# podman-hpc build --tag desihub/fastspecfit-test --file ./Containerfile-test .
+# podman-hpc migrate desihub/fastspecfit-test
+# podman-hpc run --rm -it desihub/fastspecfit-test /bin/bash
+# srun --ntasks=2 podman-hpc run --rm --mpi --entrypoint= -it desihub/fastspecfit-test /usr/local/bin/test-mpi-fastspecfit --mp=4
+
+FROM desihub/fastspecfit-base:1.0
+
+WORKDIR /src
+
+RUN ln -s "$(which python3)" /usr/bin/python
+
+ENV PIP_ROOT_USER_ACTION=ignore
+
+RUN for x in \
+    wheel \
+    setuptools \
+    pytest \
+    numpy \
+    numba \
+    mpi4py \
+    ipython \
+    ipykernel \
+    ; do pip3 install --break-system-packages $x; done \
+    && rm -Rf /root/.cache/pip
+
+COPY test-mpi-fastspecfit /usr/local/bin/test-mpi-fastspecfit
+RUN chmod +x /usr/local/bin/test-mpi-fastspecfit
+
+# set prompt and default shell
+SHELL ["/bin/bash", "-c"]
+ENTRYPOINT ["/bin/bash", "-c"]
+CMD ["/bin/bash"]
+
+LABEL Maintainer="John Moustakas [email protected]"

docker/README.md

@@ -29,7 +29,7 @@ podman-hpc push desihub/fastspecfit:3.1.1
 
 podman-hpc run --userns keep-id --rm --volume=/dvs_ro/cfs/cdirs:/dvs_ro/cfs/cdirs --volume=/pscratch/sd/i/ioannis:/scratch -it fastspecfit:3.1.1 /bin/bash
 
-
+podman-hpc run --userns keep-id --group-add keep-groups --rm --volume=/dvs_ro/cfs/cdirs:/dvs_ro/cfs/cdirs --volume=/global/cfs/cdirs:/global/cfs/cdirs --volume=/pscratch/sd/i/ioannis:/scratch --env NUMBA_CACHE_DIR=/scratch/numba_cache -it desihub/fastspecfit:3.1.1 /bin/bash
 
 List the available images:
 ```

docker/test-mpi-fastspecfit

@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+"""
+Simple test of the fastspecfit MPI wrapper.
+
+"""
+import os
+import numpy as np
+
+
+def _one_test(mpargs):
+    one_test(*mpargs)
+
+
+def one_test(rank, subrank):
+    print(f'one_test: rank {rank}, subrank {subrank}')
+
+
+def do_test(base_number, rank=0, mp=1, comm=None, nofuture=False):
+
+    print(f'do_test: rank {rank}: base number={base_number}, mp={mp}')
+
+    subranks = base_number + np.arange(mp)
+
+    if comm is not None and nofuture is False:
+        from mpi4py.futures import MPIPoolExecutor as Pool
+    else:
+        from multiprocessing import Pool
+
+    mpargs = [(rank, subrank) for subrank in subranks]
+    with Pool(mp) as P:
+        P.map(_one_test, mpargs)
+
+
+def main():
+    """Main wrapper on fastphot and fastspec.
+
+    """
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--mp', type=int, default=1, help='Number of multiprocessing processes per MPI rank or node.')
+    parser.add_argument('--nofuture', action='store_true', help='Do not use mpi4py.futures.')
+    parser.add_argument('--nompi', action='store_true', help='Do not use MPI parallelism.')
+    args = parser.parse_args()
+
+    if args.nompi:
+        comm = None
+    else:
+        try:
+            from mpi4py import MPI
+            comm = MPI.COMM_WORLD
+        except ImportError:
+            comm = None
+
+    if comm is None:
+        rank, size = 0, 1
+    else:
+        rank, size = comm.rank, comm.size
+
+    print(f'main: I am on rank {rank}/{size}')
+
+    # Main rank is responsible for planning.
+    if rank == 0:
+        base_number = np.arange(size) * 128
+
+        # https://docs.nersc.gov/development/languages/python/parallel-python/#use-the-spawn-start-method
+        if args.nofuture and args.mp > 1 and 'NERSC_HOST' in os.environ:
+            import multiprocessing
+            multiprocessing.set_start_method('spawn')
+    else:
+        base_number = []
+
+    if comm:
+        base_number = comm.bcast(base_number, root=0)
+
+    do_test(base_number[rank], rank=rank, mp=args.mp,
+            comm=comm, nofuture=args.nofuture)
+
+
+if __name__ == '__main__':
+    main()