Make sure shm.dll can find mkl.dll in torch 2.3+

metatensor · Nov 29, 2024 · cb5e1fd · cb5e1fd
1 parent 1fc2ffc
commit cb5e1fd
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 10 deletions.
diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
@@ -95,7 +95,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-wheel-${{ matrix.os }}-${{ matrix.cibw-arch }}
+          name: wheel-${{ matrix.os }}-${{ matrix.cibw-arch }}
           path: ./wheelhouse/*.whl
 
   build-torch-wheels:
@@ -214,13 +214,13 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-torch-single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
+          name: torch-single-version-wheel-${{ matrix.torch-version }}-${{ matrix.os }}-${{ matrix.arch }}
           path: ./wheelhouse/*.whl
 
   merge-torch-wheels:
     needs: build-torch-wheels
     runs-on: ubuntu-22.04
-    name: merge metatensor-torch ${{ matrix.name }}
+    name: merge featomic-torch ${{ matrix.name }}
     strategy:
       matrix:
         include:
@@ -287,10 +287,10 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: featomic-torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
+          name: torch-wheel-${{ matrix.os }}-${{ matrix.arch }}
           path: ./wheelhouse/*.whl
 
-  build-others:
+  build-sdist:
     runs-on: ubuntu-22.04
     name: sdist
     steps:
@@ -319,7 +319,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: others
+          name: sdist
           path: |
             dist/*.tar.gz
             dist/cxx/*.tar.gz
@@ -336,7 +336,7 @@ jobs:
 
   merge-and-release:
     name: Merge and release wheels/sdists
-    needs: [build-wheels, merge-torch-wheels, build-others]
+    needs: [build-wheels, merge-torch-wheels, build-sdist]
     runs-on: ubuntu-22.04
     permissions:
       contents: write
@@ -345,21 +345,21 @@ jobs:
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          pattern: featomic-wheel-*
+          pattern: wheel-*
           merge-multiple: true
 
       - name: Download featomic-torch wheels
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          pattern: featomic-torch-wheel-*
+          pattern: torch-wheel-*
           merge-multiple: true
 
       - name: Download sdists
         uses: actions/download-artifact@v4
         with:
           path: wheels
-          name: others
+          name: sdist
 
       - name: Re-upload a single wheels artifact
         uses: actions/upload-artifact@v4

diff --git a/python/featomic_torch/setup.py b/python/featomic_torch/setup.py
@@ -268,6 +268,38 @@ def create_version_number(version):
 
 
 if __name__ == "__main__":
+    if sys.platform == "win32":
+        # On Windows, starting with PyTorch 2.3, the file shm.dll in torch has a
+        # dependency on mkl DLLs. When building the code using pip build isolation, pip
+        # installs the mkl package in a place where the os is not trying to load
+        #
+        # This is a very similar fix to https://github.com/pytorch/pytorch/pull/126095,
+        # except only applying when importing torch from a build-isolation virtual
+        # environment created by pip (`python -m build` does not seems to suffer from
+        # this).
+        import wheel
+
+        pip_virtualenv = os.path.realpath(
+            os.path.join(
+                os.path.dirname(wheel.__file__),
+                "..",
+                "..",
+                "..",
+                "..",
+            )
+        )
+        mkl_dll_dir = os.path.join(
+            pip_virtualenv,
+            "normal",
+            "Library",
+            "bin",
+        )
+
+        if os.path.exists(mkl_dll_dir):
+            os.add_dll_directory(mkl_dll_dir)
+
+        # End of Windows/MKL/PIP hack
+
     if not os.path.exists(FEATOMIC_TORCH_SRC):
         # we are building from a sdist, which should include featomic-torch
         # sources as a tarball