[MLIR][Python] Fix stubgen/PYTHONPATH collision/bug

[makslevental]

If PYTHONPATH is set and points to the build location of the python bindings package then when stubgen runs, _mlir will get imported twice and bad things will happen (e.g., Assertion !instance && “PyGlobals already constructed”’). This happens because mlir is a namespace package and the importer/loader can't distinguish between
mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD. Or something like that. The fix is to filter out any entries in PYTHONPATH that end in MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.. (e.g., python_packages/mlir_core/).

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

If PYTHONPATH is set and points to the build location of the python bindings package then when stubgen runs, _mlir will get imported twice and bad things will happen (e.g., Assertion !instance && “PyGlobals already constructed”’). This happens because mlir is a namespace package and so importer/loader can't distinguish between
mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD. Or something like that. The fix is to filter out any entries in PYTHONPATH that end in MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.. (e.g., python_packages/mlir_core/).

---

Full diff: https://github.com/llvm/llvm-project/pull/161307.diff

1 Files Affected:

• (modified) mlir/cmake/modules/AddMLIRPython.cmake (+21-1)

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 208cbdd1dd535..1127f9410e488 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -173,9 +173,29 @@ function(mlir_generate_type_stubs)
   if(ARG_VERBOSE)
     message(STATUS "Generating type-stubs outputs ${_generated_type_stubs}")
   endif()
+
+  # If PYTHONPATH is set and points to the build location of the python package then when stubgen runs, _mlir will get
+  # imported twice and bad things will happen (e.g., Assertion `!instance && “PyGlobals already constructed”’).
+  # This happens because mlir is a namespace package and so importer/loader can't distinguish between
+  # mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD.
+  # So try to filter out any entries in PYTHONPATH that end in "MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.."
+  # (e.g., python_packages/mlir_core/).
+  set(_pythonpath "$ENV{PYTHONPATH}")
+  cmake_path(SET _install_prefix NORMALIZE "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/..")
+  string(REGEX REPLACE "/$" "" _install_prefix "${_install_prefix}")
+  if(WIN32)
+    set(_path_sep ";")
+  else()
+    set(_path_sep ":")
+    # `;` is the CMake list delimiter so Windows paths are automatically lists
+    # but Unix paths can be made into lists by replacing `:` with `;`
+    string(REPLACE "${_path_sep}" ";" _pythonpath "${_pythonpath}")
+  endif()
+  list(FILTER _pythonpath EXCLUDE REGEX "(${_install_prefix}|${_install_prefix}/)$")
+  string(JOIN "${_path_sep}" _pythonpath ${_pythonpath})
   add_custom_command(
     OUTPUT ${_generated_type_stubs}
-    COMMAND ${_nb_stubgen_cmd}
+    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${_pythonpath} ${_nb_stubgen_cmd}
     WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
     DEPENDS "${ARG_DEPENDS_TARGETS}"
     DEPFILE "${_depfile}"

[PragmaTwice]

Tried on my local and it did fix the issue. Thanks!