[PJRT] Fix ninja not found error while re-building the plugin (#19553)

If we want to build the PJRT CPU plugin, we'll run something like
```
pip install --no-deps -v ./integrations/pjrt/python_packages/iree_cpu_plugin/
```

It works well in the first run. But if we did some changes, and want to
run it the second time, some errors will appear: cmake cannot find the
ninja in the first run anymore because it's in a temp build environment
and removed after the first build is finished.

We can remove the build dir to solve this problem. But it will cause a
full rebuild, and is quite annoying : )

Since IREE compiler doesn't have such issue so I checked its build
script, and I found that it's solved via the function
`maybe_nuke_cmake_cache` in its
[setup.py](https://github.com/iree-org/iree/blob/76a7b893e4c62d52eae2c165bdb23952a8589689/compiler/setup.py#L177).
So I copy it into setup.py of the PJRT plugin and did some modification:
- I think the PJRT plugin doesn't rely on CPython API (although it
builds a shared library) so we don't need to pin Python version;
- the build dir should be passed via a parameter since we have plugins
for different platforms (cpu/cuda/rocm..).

Also, I used this chance to add `cmake` to build dependencies, in case
of some users don't have cmake installed in the system.

ci-exactly: build_packages, test_pjrt

Signed-off-by: PragmaTwice <[email protected]>

Author: PragmaTwice

integrations/pjrt/python_packages/_setup_support/iree_pjrt_setup.py

@@ -121,6 +121,42 @@ def run(self):
         self.run_command("build_py")
 
 
+def maybe_nuke_cmake_cache(cmake_build_dir):
+    # From run to run under pip, we can end up with different paths to ninja,
+    # which isn't great and will confuse cmake. Detect if the location of
+    # ninja changes and force a cache flush.
+    ninja_path = ""
+    try:
+        import ninja
+    except ModuleNotFoundError:
+        pass
+    else:
+        ninja_path = ninja.__file__
+    expected_stamp_contents = f"{ninja_path}"
+
+    # In order to speed things up on CI and not rebuild everything, we nuke
+    # the CMakeCache.txt file if the path to ninja changed.
+    # Ideally, CMake would let us reconfigure this dynamically... but it does
+    # not (and gets very confused).
+    NINJA_STAMP_FILE = os.path.join(cmake_build_dir, "ninja_stamp.txt")
+    if os.path.exists(NINJA_STAMP_FILE):
+        with open(NINJA_STAMP_FILE, "rt") as f:
+            actual_stamp_contents = f.read()
+            if actual_stamp_contents == expected_stamp_contents:
+                # All good.
+                return
+
+    # Mismatch or not found. Clean it.
+    cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
+    if os.path.exists(cmake_cache_file):
+        print("Removing CMakeCache.txt because Ninja path changed", file=sys.stderr)
+        os.remove(cmake_cache_file)
+
+    # And write.
+    with open(NINJA_STAMP_FILE, "wt") as f:
+        f.write(expected_stamp_contents)
+
+
 class BaseCMakeBuildPy(_build_py):
     def run(self):
         self.build_default_configuration()
@@ -131,6 +167,7 @@ def run(self):
 
     def build_configuration(self, cmake_build_dir, extra_cmake_args=()):
         subprocess.check_call(["cmake", "--version"])
+        maybe_nuke_cmake_cache(cmake_build_dir)
 
         cfg = os.getenv("IREE_CMAKE_BUILD_TYPE", "Release")
 

integrations/pjrt/python_packages/iree_cpu_plugin/pyproject.toml

@@ -3,5 +3,6 @@ requires = [
     "setuptools>=42",
     "wheel",
     "ninja",
+    "cmake",
 ]
 build-backend = "setuptools.build_meta"

integrations/pjrt/python_packages/iree_cuda_plugin/pyproject.toml

@@ -3,5 +3,6 @@ requires = [
     "setuptools>=42",
     "wheel",
     "ninja",
+    "cmake",
 ]
 build-backend = "setuptools.build_meta"

integrations/pjrt/python_packages/iree_rocm_plugin/pyproject.toml

@@ -3,5 +3,6 @@ requires = [
     "setuptools>=42",
     "wheel",
     "ninja",
+    "cmake",
 ]
 build-backend = "setuptools.build_meta"

integrations/pjrt/python_packages/iree_vulkan_plugin/pyproject.toml

@@ -3,5 +3,6 @@ requires = [
     "setuptools>=42",
     "wheel",
     "ninja",
+    "cmake",
 ]
 build-backend = "setuptools.build_meta"