Merge pull request #71 from trustimaging/save-bounds-hotfix

Fix compatibility issue with torch

docker/jupyter.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
 source ~/.bashrc
-jupyter notebook --ip=0.0.0.0 --port=8888 --allow-root
+imrun notebook --ip=0.0.0.0 --port=8888 --allow-root

environment.yml

@@ -25,22 +25,26 @@ dependencies:
   - pyflakes
   - pytest
   - pytest-cov
+  - pytest-runner
+  - pytorch
   - pyyaml>=5.4
   - pyqt5-sip
-  - pytest-runner
   - pyzmq>=22.1
   - scikit-image
   - scipy>=1.6
   - sphinx
   - sphinx_rtd_theme
   - tornado>=6.1
+  - tqdm
   - traitsui
+  - ipympl
   - uvloop
   - zlib
   - zict>=3.0.0
   - pip:
       - devito>=4.6
+      - ndsplines
       - python-daemon
       - py-libnuma
-      - tblib
       - resampy
+      - tblib

mosaic/cli/findomp.py

@@ -0,0 +1,48 @@
+
+import re
+import os
+import sys
+
+
+def go():
+    # Get CUDA compilers path
+    comp_path = None
+
+    # Try to get from environment variables
+    envs = ['CUDA_HOME', 'NVHPC_CUDA_HOME', 'CUDA_ROOT']
+    for env in envs:
+        cuda_home = os.environ.get(env, None)
+        if cuda_home:
+            comp_path = os.path.join(os.path.join(os.path.dirname(cuda_home), 'compilers'), 'lib')
+            break
+
+    if comp_path is None:
+        hpcsdk_home = os.environ.get('HPCSDK_HOME')
+        if hpcsdk_home:
+            comp_path = os.path.join(os.path.join(hpcsdk_home, 'compilers'), 'lib')
+
+    # If not, try to get from LD_LIBRARY_PATH
+    if comp_path is None:
+        library_path = os.environ.get('LD_LIBRARY_PATH', '')
+        for path in library_path.split(':'):
+            if re.match('.*/nvidia/hpc_sdk/.*/?compilers/lib', path) or \
+                    re.match('.*/nvhpc/.*/?compilers/lib', path):
+                comp_path = path
+
+    if comp_path is None or not os.path.exists(comp_path):
+        return ''
+
+    # Now try to get libgomp path
+    libnames = ['libgomp', 'libiomp', 'libomp']
+    extnames = ['.so', '.dylib', '.dll']
+    for lib in libnames:
+        for ext in extnames:
+            lib_path = os.path.join(comp_path, lib + ext)
+            if os.path.exists(lib_path):
+                return lib_path
+
+    return ''
+
+
+if __name__ == '__main__':
+    sys.exit(go())

mosaic/cli/imrun

@@ -0,0 +1,3 @@
+#!/bin/bash
+preload=$((findomp) 2>&1)
+LD_PRELOAD="$preload" jupyter "$@"

mosaic/cli/mrun

@@ -0,0 +1,3 @@
+#!/bin/bash
+preload=$((findomp) 2>&1)
+LD_PRELOAD="$preload" mrun_ "$@"

requirements.txt

@@ -8,9 +8,11 @@ devito
 flake8
 gputil
 h5py>=3.2
+multiprocess
 nbval
+ndsplines
 numpy>=1.20
-multiprocess
+notebook>=6.4
 pre_commit
 psutil
 pyflakes
@@ -26,6 +28,10 @@ scikit-image
 scipy>=1.6
 sphinx
 sphinx_rtd_theme
+tornado>=6.1
 tblib
+torch
+tqdm
+ipympl
 uvloop
 zict>=3.0.0

setup.py

@@ -47,11 +47,17 @@
         Extension('_profile',
                   sources=['mosaic/profile/_profile.c'])
     ],
+    scripts=[
+        'mosaic/cli/mrun',
+        'mosaic/cli/imrun',
+    ],
     entry_points={
         'console_scripts': [
             'mrun=mosaic.cli.mrun:go',
+            'mrun_=mosaic.cli.mrun:go',
             'mscript=mosaic.cli.mscript:go',
             'mprof=mosaic.cli.mprof:go',
+            'findomp=mosaic.cli.findomp:go',
         ]
     },
     zip_safe=False,

stride/physics/common/devito.py

@@ -509,7 +509,9 @@ def _time_undersampled(self, name, factor, bounds=None, offset=None):
                                                  factor=factor,
                                                  condition=condition)
 
-        buffer_size = (bounds[1] - bounds[0] + factor) // factor + 1
+        # buffer_size = (bounds[1] - bounds[0] + factor) // factor + 1
+        # TODO Force larger buffer size to prevent devito issue
+        buffer_size = (self.time.extended_num - 1 - 0 + factor) // factor + 1
 
         return time_under, buffer_size