Release 0.10

.github/workflows/container.yml

@@ -0,0 +1,53 @@
+name: Build and Push Container Image
+
+on:
+  push:
+    tags:
+      - v*
+
+jobs:
+  build-and-push-image:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    strategy:
+      matrix:
+        platform: [full, cupy]
+        mpi: [openmpi, mpich]
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract Version from Tag
+        id: tags
+        run: echo version=$(echo "${{ github.ref_name }}" | cut -c 2-) >> $GITHUB_OUTPUT
+
+      - name: Docker Metadata
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}_${{ matrix.platform }}_${{ matrix.mpi }}
+          tags: |
+            type=raw,value=${{ steps.tags.outputs.version}}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build and Push Image
+        uses: docker/build-push-action@v6
+        with:
+          push: true
+          context: .
+          tags: ${{ steps.meta.outputs.tags }}
+          target: runtime
+          build-args: |
+            PLATFORM=${{ matrix.platform}}
+            MPI=${{ matrix.mpi }}

.github/workflows/test.yml

@@ -51,7 +51,7 @@ jobs:
         miniforge-version: latest
     - name: Install ${{ matrix.conda-env }} dependencies
       run: |
-        sed -i 's/python/python=${{ matrix.python-version }}/' dependencies_${{ matrix.conda-env }}.yml
+        sed -i '/- python/d' dependencies_${{ matrix.conda-env }}.yml
         conda install -c conda-forge mpich
         conda env update --file dependencies_${{ matrix.conda-env }}.yml --name ptypy_env
         conda install -c conda-forge flake8 pytest pytest-cov

Dockerfile

@@ -5,69 +5,92 @@ ARG MPI=openmpi
 ARG PLATFORM=cupy
 
 # Select CUDA version
-ARG CUDAVERSION=12.4
+ARG CUDAVERSION=12.8
 
 # Pull from mambaforge and install XML and ssh
-FROM condaforge/mambaforge as base
+FROM condaforge/mambaforge AS base
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && apt-get install -y libxml2 ssh
 
 # Pull from base image and install OpenMPI/MPICH
-FROM base as mpi
+FROM base AS mpi
 ARG MPI
 RUN mamba install -n base -c conda-forge ${MPI}
 
-# Pull from MPI build install core dependencies
-FROM base as core
-COPY ./dependencies_core.yml ./dependencies.yml
-RUN mamba env update -n base -f dependencies.yml
+# Pull from MPI build and install core dependencies
+FROM base AS core
+RUN mamba install -n base -y -c conda-forge \
+    python numpy scipy h5py pip
 
 # Pull from MPI build and install full dependencies
-FROM mpi as full
-COPY ./dependencies_full.yml ./dependencies.yml
-RUN mamba env update -n base -f dependencies.yml
+FROM mpi AS full
+ARG MPI
+RUN mamba install -n base -y -c conda-forge \
+    python numpy scipy matplotlib h5py \ 
+    pyzmq mpi4py[build=*${MPI}*] packaging \
+    pillow pyfftw pyyaml pip
 
 # Pull from MPI build and install accelerate/pycuda dependencies
-FROM mpi as pycuda
-ARG CUDAVERSION
-COPY ./ptypy/accelerate/cuda_pycuda/dependencies.yml ./dependencies.yml
-COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba install cuda-version=${CUDAVERSION} && \
-    mamba env update -n base -f dependencies.yml && \
-    mamba env update -n base -f dependencies_cufft.yml
+FROM mpi AS pycuda
+ARG CUDAVERSION MPI
+RUN mamba install -n base -y -c conda-forge -c nvidia \
+    python numpy scipy matplotlib h5py pyzmq mpi4py[build=*${MPI}*] \
+    pillow pyfftw pyyaml compilers pip \
+    reikna pycuda cuda-nvcc cuda-cudart-dev cuda-version=${CUDAVERSION}
 
 # Pull from MPI build and install accelerate/cupy dependencies
-FROM mpi as cupy
-ARG CUDAVERSION
-COPY ./ptypy/accelerate/cuda_cupy/dependencies.yml ./dependencies.yml
-COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba install cuda-version=${CUDAVERSION} && \
-    mamba env update -n base -f dependencies.yml && \
-    mamba env update -n base -f dependencies_cufft.yml
+FROM mpi AS cupy
+ARG CUDAVERSION MPI
+RUN mamba install -n base -y -c conda-forge \
+    python numpy scipy matplotlib h5py pyzmq mpi4py[build=*${MPI}*] \
+    pillow pyfftw pyyaml compilers pip \
+    cupy cuda-version=${CUDAVERSION}
+RUN mamba clean -y -a
 
 # Pull from platform specific image and install ptypy 
-FROM ${PLATFORM} as build
+FROM ${PLATFORM} AS build
 COPY pyproject.toml ./
 COPY ./templates ./templates
 COPY ./benchmark ./benchmark
-COPY ./cufft ./cufft
 COPY ./ptypy ./ptypy
 RUN pip install .
 
-# For core/full build, no post processing needed
-FROM build as core-post
-FROM build as full-post
+# For core build, clean up conda env
+FROM build AS core-post
+RUN mamba clean -y -a
+
+# For full build, clean up conda env
+FROM build AS full-post
+RUN mamba clean -y -a
 
 # For pycuda build, install filtered cufft
-FROM build as pycuda-post
+FROM build AS pycuda-post
+ARG CUDAVERSION
+RUN mamba install -n base -y -c conda-forge -c nvidia \
+    python cmake>=3.8.0 pybind11 compilers \
+    cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static cuda-version=${CUDAVERSION}
+COPY ./cufft ./cufft
 RUN pip install ./cufft
+RUN mamba remove -n base -y \
+    cmake pybind11 cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static
+RUN mamba clean -y -a
 
-# For pycuda build, install filtered cufft
-FROM build as cupy-post
+# For cupy build, install filtered cufft
+FROM build AS cupy-post
+ARG CUDAVERSION
+RUN mamba install -n base -y -c conda-forge -c nvidia \
+    python cmake>=3.8.0 pybind11 compilers \
+    cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static cuda-version=${CUDAVERSION}
+COPY ./cufft ./cufft
 RUN pip install ./cufft
+RUN mamba remove -n base -y \
+    cmake pybind11 libcufft-dev libcufft-static
+RUN mamba clean -y -a
 
 # Platform specific runtime container
-FROM ${PLATFORM}-post as runtime
+FROM ${PLATFORM}-post AS runtime
+RUN useradd --user-group ptypy-user
+USER ptypy-user
 
 # Run PtyPy run script as entrypoint
 ENTRYPOINT ["ptypy.cli"]

cufft/extensions.py

@@ -147,7 +147,6 @@ def build_extension(self, ext):
         if has_cu:
             old_compiler = self.compiler
             self.compiler = NvccCompiler(verbose=old_compiler.verbose,
-                                        dry_run=old_compiler.dry_run,
                                         force=old_compiler.force) # this is our bespoke compiler
             super(CustomBuildExt, self).build_extension(ext)
             self.compiler=old_compiler

ptypy/accelerate/cuda_cupy/engines/ML_cupy.py

@@ -14,6 +14,7 @@
 import numpy as np
 import cupy as cp
 import cupyx
+import os
 
 from ptypy.engines import register
 from ptypy.accelerate.base.engines.ML_serial import ML_serial, BaseModelSerial
@@ -35,6 +36,11 @@
 MAX_BLOCKS = 99999
 # MAX_BLOCKS = 3  # can be used to limit the number of blocks, simulating that they don't fit
 
+# Estimate the device memory safety margin as fraction of total device memory
+device_memory_fractional_safety_margin = 0.025
+if "PTYPY_DEVICE_MEM_SAFETY" in os.environ:
+    device_memory_fractional_safety_margin = float(os.environ["PTYPY_DEVICE_MEM_SAFETY"])
+
 
 @register()
 class ML_cupy(ML_serial):
@@ -160,9 +166,11 @@ def _setup_kernels(self):
         # as both will be used for allocations
         mempool = cp.get_default_memory_pool()
         mem = cp.cuda.runtime.memGetInfo()[0] + mempool.total_bytes() - mempool.used_bytes()
+        tot = cp.cuda.runtime.memGetInfo()[1]
+        safety_margin = max(device_memory_fractional_safety_margin * tot, 400 * 1024 * 1024)
 
-        # leave 200MB room for safety
-        fit = int(mem - 200 * 1024 * 1024) // blk
+        # leave room for safety
+        fit = int(mem - safety_margin) // blk
         if not fit:
             log(1, "Cannot fit memory into device, if possible reduce frames per block. Exiting...")
             raise SystemExit("ptypy has been exited.")

ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py

@@ -13,6 +13,7 @@
     :license: see LICENSE for details.
 """
 
+import os
 import numpy as np
 import cupy as cp
 import cupyx
@@ -32,6 +33,11 @@
 MAX_BLOCKS = 99999
 # MAX_BLOCKS = 3  # can be used to limit the number of blocks, simulating that they don't fit
 
+# Estimate the device memory safety margin as fraction of total device memory
+device_memory_fractional_safety_margin = 0.025
+if "PTYPY_DEVICE_MEM_SAFETY" in os.environ:
+    device_memory_fractional_safety_margin = float(os.environ["PTYPY_DEVICE_MEM_SAFETY"])
+
 __all__ = ['DM_cupy_stream', 'RAAR_cupy_stream']
 
 
@@ -65,9 +71,11 @@ def _setup_kernels(self):
         # as both will be used for allocations
         mempool = cp.get_default_memory_pool()
         mem = cp.cuda.runtime.memGetInfo()[0] + mempool.total_bytes() - mempool.used_bytes()
-
-        # leave 200MB room for safety
-        fit = int(mem - 200 * 1024 * 1024) // blk
+        tot = cp.cuda.runtime.memGetInfo()[1]
+        safety_margin = max(device_memory_fractional_safety_margin * tot, 400 * 1024 * 1024)
+
+        # leave room for safety
+        fit = int(mem - safety_margin) // blk
         if not fit:
             log(1, "Cannot fit memory into device, if possible reduce frames per block. Exiting...")
             raise SystemExit("ptypy has been exited.")

ptypy/cli/command_line_interface.py

@@ -31,6 +31,8 @@ def parse():
                             help="Provide parameter configuration as a JSON string.")
     parser.add_argument('--output-folder', '-o', type=str,
                         help="The path we want the outputs to exist in (will get created).")
+    parser.add_argument('--use-timestamp', action="store_true",
+                        help="Turn on timestamping of output folder")
     parser.add_argument('--ptypy-level', '-l', default=5, type=str,
                         help="The level we want to run to ptypy to.")
     parser.add_argument('--identifier', '-i', type=str, default=None,
@@ -57,9 +59,9 @@ def run(args):
 
     # Load parameter tree from file or JSON string
     if args.file:
-        p = create_parameter_tree(load_config_as_dict_from_file(args.file))
+        p, extra = create_parameter_tree(load_config_as_dict_from_file(args.file))
     if args.json:
-        p = create_parameter_tree(json.loads(args.json))
+        p, extra = create_parameter_tree(json.loads(args.json))
     p.run = args.identifier
 
     # TODO
@@ -74,7 +76,7 @@ def run(args):
     # 
     if args.output_folder is not None:
         p.io.home = get_output_file_name(args)
-        p.io.rfile = "%s.ptyr" % get_output_file_name(args)
+        p.io.rfile = "%(run)s_%(engine)s_%(iterations)04d.ptyr"
         #parameters.io.autosave = u.Param(active=True)
         #log(3, "Autosave is on, with io going in {}, and the final reconstruction into {}".format(parameters.io.home,
         #                                                                                        parameters.io.rfile))
@@ -83,8 +85,8 @@ def run(args):
         p.io.autosave = u.Param(active=False)
         log("info", "Autosave is off. No output will be saved.")
 
-    # Substitute %(run) with in ptyscan 
-    substitute_id_in_ptyscan(p)
+    # Substitute %(run) with in ptyscan
+    substitute_id_in_ptyscan(p, extra)
 
     # Run PtyPy to given level
     P = Ptycho(p, level=args.ptypy_level)
@@ -110,9 +112,16 @@ def create_parameter_tree(params) -> u.Param:
         parameters_to_run.update(previous_parameters)
     if params['parameter_tree'] is not None:
         parameters_to_run.update(params['parameter_tree'], Convert=True)
-    return parameters_to_run
 
-def substitute_id_in_ptyscan(params):
+    # Additional non-ptypy params
+    extra = {}
+    for k,v in params.items():
+        if k != "parameter_tree":
+            extra[k] = v
+
+    return parameters_to_run, extra
+
+def substitute_id_in_ptyscan(params, extra):
     def _substitute(d, p):
         for k, v in d.items():
             if isinstance(v, MutableMapping):
@@ -122,15 +131,16 @@ def _substitute(d, p):
                 d[k] = v % p
     for scan_key, scan in params.scans.items():
         data_entry = scan.data
-        _substitute(data_entry, params)
+        _substitute(data_entry, params | extra)
 
 def get_output_file_name(args):
     from datetime import datetime
     now = datetime.now()
+    output_path = "{}/scan".format(args.output_folder)
     if args.identifier is not None:
-        output_path = "{}/scan_{}_{}".format(args.output_folder, args.identifier, now.strftime("%Y%m%d%H%M%S"))
-    else:
-        output_path = "{}/scan_{}".format(args.output_folder, now.strftime("%Y%m%d%H%M%S"))
+        output_path += "_{}".format(args.identifier)
+    if args.use_timestamp:
+        output_path += "_{}".format(now.strftime("%Y%m%d%_H%M%S"))
     log("info", "Output is going in: {}".format(output_path))
     return output_path