bug fix for memory chunks in mpcomplex, github actions again

NYU-DiffusionMRI · Jun 14, 2024 · 56aeb08 · 56aeb08
1 parent 747815c
commit 56aeb08
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 23 deletions.
diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml
@@ -29,12 +29,14 @@ jobs:
         run: |
           docker run --rm -v $GITHUB_WORKSPACE:/project quay.io/pypa/manylinux2014_${{ matrix.architecture }} /bin/bash -c "
           yum install -y fftw-devel && \
-          for PYBIN in /opt/python/cp${{ matrix.python-version }}*/bin; do
-            \$PYBIN/python -m ensurepip && \
-            \$PYBIN/pip install -U pip setuptools wheel auditwheel pybind11 && \
-            cd /project && \
-            \$PYBIN/python setup.py bdist_wheel && \
-            \$PYBIN/auditwheel repair dist/*.whl -w dist/
+          for PYBIN in /opt/python/*${{ matrix.python-version }}*/bin; do
+            if [ -d \$PYBIN ]; then
+              \$PYBIN/python -m ensurepip && \
+              \$PYBIN/pip install -U pip setuptools wheel auditwheel pybind11 && \
+              cd /project && \
+              \$PYBIN/python setup.py bdist_wheel && \
+              \$PYBIN/auditwheel repair dist/*.whl -w dist/
+            fi
           done
           "
       - name: Upload to PyPI
@@ -62,16 +64,21 @@ jobs:
           brew install fftw
       - name: Build wheels
         run: |
-          if [[ "$(uname -m)" == "arm64" ]]; then
+          if [[ "${{ matrix.architecture }}" == "arm64" ]]; then
             arch -arm64 python setup.py bdist_wheel
           else
             python setup.py bdist_wheel
           fi
       - name: Repair wheels
         run: |
           pip install delocate
-          delocate-listdeps dist/*.whl  # lists library dependencies
-          delocate-wheel dist/*.whl  # copies library dependencies into the wheel
+          if [[ "${{ matrix.architecture }}" == "arm64" ]]; then
+            arch -arm64 delocate-listdeps dist/*.whl  # lists library dependencies
+            arch -arm64 delocate-wheel dist/*.whl  # copies library dependencies into the wheel
+          else
+            delocate-listdeps dist/*.whl  # lists library dependencies
+            delocate-wheel dist/*.whl  # copies library dependencies into the wheel
+          fi
       - name: Upload to PyPI
         uses: pypa/gh-action-pypi-publish@master
         with:
@@ -95,8 +102,16 @@ jobs:
           python -m pip install --upgrade pip setuptools wheel pybind11
       - name: Download and install FFTW
         run: |
-          curl -L -o fftw.zip https://www.fftw.org/fftw-3.3.10-dll64.zip
-          unzip fftw.zip -d fftw
+          curl -L -o fftw.zip http://www.fftw.org/fftw-3.3.10-dll64.zip
+          if (-not (Test-Path fftw.zip)) {
+            Write-Error "Failed to download fftw.zip"
+            exit 1
+          }
+          Expand-Archive fftw.zip -DestinationPath fftw
+          if (-not (Test-Path fftw)) {
+            Write-Error "Failed to unzip fftw.zip"
+            exit 1
+          }
           copy fftw/*.dll C:\Windows\System32
       - name: Build wheels
         run: |

diff --git a/lib/mpcomplex.py b/lib/mpcomplex.py
@@ -236,19 +236,14 @@ def im_reconstruct(self, wp, image):
             rec_img = self.unpad(rec_img, self.pwidth[:][:-1])
         elif image.ndim == 3:
             # Process in chunks to manage memory usage
-            chunk_size = max(1, image.shape[2] // 10)
-            rec_img_chunks = []
-            for i in range(0, image.shape[2], chunk_size):
-                end = min(i + chunk_size, image.shape[2])
-                S = Parallel(n_jobs=self.n_cores, prefer='threads')\
-                    (delayed(self.patchav)(
-                        wp, image[:, :, j].flatten(), self.temp.flatten()
-                        ) for j in range(i, end))
-                rec_img_chunks.append(np.array(S).transpose(1, 2, 0))
-                del S
-                gc.collect()
-            rec_img = np.concatenate(rec_img_chunks, axis=2)
+            S = Parallel(n_jobs=-3, prefer='threads')\
+            (delayed(self.patchav)(
+                wp, image[:,:,i].flatten(), self.temp.flatten()
+                ) for i in range(self.imsize[-1]))
+            rec_img = np.array(S).transpose(1,2,3,0)
             rec_img = self.unpad(rec_img, self.pwidth)
+
+        gc.collect()
         return rec_img
 
     def process(self):