bugfix: Ninja race condition fix (#2339)

<!-- .github/pull_request_template.md -->

## 📌 Description

**Fix race condition in JIT compilation for multi-GPU/multi-process
environments**

This PR resolves a race condition in multi-process environments where
concurrent Ninja builds contended for a shared `.ninja_log` file
(`ninja: error: opening build log: No such file or directory`), causing
intermittent failures when multiple processes compile different
FlashInfer modules simultaneously.


Key changes:
1. Isolated workdirs for runtime JIT: Each module now builds in its own
subdirectory (`cached_ops/<module>/`), isolating .ninja_log files
between concurrent builds
2. Absolute output paths: Ninja build files use absolute paths for
object and output files, ensuring correct output locations regardless of
workdir
3. Preserved AOT parallelism: Batch builds retain the `subninja`
approach, maintaining full parallel compilation



Aspect | Before | After
-- | -- | --
Runtime JIT workdir | Shared `cached_ops/` |
Isolated `cached_ops/<module>/`
AOT batch workdir | `cached_ops/` | `cached_ops/` (unchanged)
`.ninja_log` conflict | Yes (race condition) | No (different workdir
levels)
Ninja output paths | Relative (`$name/$name.so`) | Absolute
(`/path/to/cached_ops/<module>/<module>.so`)
Output locations | `cached_ops/<module>/<module>.so` |
`cached_ops/<module>/<module>.so` (unchanged)
AOT parallelism | Full (`subninja`) | Full (`subninja`) - preserved


<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

#2338

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Build system now isolates per-spec build outputs into dedicated build
directories and runs per-spec compilation with top-level artifact
targets to avoid cross-spec interference and race conditions.
* **New Features**
* Added a public per-spec build_dir property to configure where each
specification’s build outputs are placed.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: bkryu

flashinfer/jit/core.py

@@ -227,6 +227,10 @@ class JitSpec:
     def ninja_path(self) -> Path:
         return jit_env.FLASHINFER_JIT_DIR / self.name / "build.ninja"
 
+    @property
+    def build_dir(self) -> Path:
+        return jit_env.FLASHINFER_JIT_DIR / self.name
+
     @property
     def jit_library_path(self) -> Path:
         return jit_env.FLASHINFER_JIT_DIR / self.name / f"{self.name}.so"
@@ -238,7 +242,7 @@ def get_library_path(self) -> Path:
 
     def get_object_paths(self) -> List[Path]:
         object_paths = []
-        jit_dir = self.jit_library_path.parent
+        jit_dir = self.build_dir
         for source in self.sources:
             is_cuda = source.suffix == ".cu"
             object_suffix = ".cuda.o" if is_cuda else ".o"
@@ -264,7 +268,7 @@ def lock_path(self) -> Path:
 
     def write_ninja(self) -> None:
         ninja_path = self.ninja_path
-        ninja_path.parent.mkdir(parents=True, exist_ok=True)
+        self.build_dir.mkdir(parents=True, exist_ok=True)
         content = generate_ninja_build_for_op(
             name=self.name,
             sources=self.sources,
@@ -295,7 +299,7 @@ def build(self, verbose: bool, need_lock: bool = True) -> None:
             # Write ninja file if it doesn't exist (deferred case)
             if not self.is_ninja_generated:
                 self.write_ninja()
-            run_ninja(jit_env.FLASHINFER_JIT_DIR, self.ninja_path, verbose)
+            run_ninja(self.build_dir, self.ninja_path, verbose)
 
     def load(self, so_path: Path):
         return tvm_ffi.load_module(str(so_path))
@@ -362,7 +366,7 @@ def expand_flags(
         nvcc = os.environ.get("FLASHINFER_NVCC", f"{cuda_home}/bin/nvcc")
 
         # Build directory
-        build_dir = str(self.jit_library_path.parent.resolve())
+        build_dir = str(self.build_dir.resolve())
 
         # Generate entries for each source file
         compile_commands = []

flashinfer/jit/cpp_ext.py

@@ -276,20 +276,26 @@ def generate_ninja_build_for_op(
             ]
         )
 
+    # Use absolute paths for outputs so ninja files work with any workdir
+    # This enables isolated workdirs for runtime JIT (avoiding .ninja_log races)
+    # while still supporting subninja for parallel AOT builds
+    output_dir = jit_env.FLASHINFER_JIT_DIR / name
+
     objects = []
     for source in sources:
         is_cuda = source.suffix == ".cu"
         object_suffix = ".cuda.o" if is_cuda else ".o"
         cmd = "cuda_compile" if is_cuda else "compile"
         obj_name = source.with_suffix(object_suffix).name
-        obj = f"$name/{obj_name}"
+        obj = str((output_dir / obj_name).resolve())
         objects.append(obj)
         lines.append(f"build {obj}: {cmd} {source.resolve()}")
 
     lines.append("")
     link_rule = "nvcc_link" if needs_device_linking else "link"
-    lines.append(f"build $name/$name.so: {link_rule} " + " ".join(objects))
-    lines.append("default $name/$name.so")
+    output_so = str((output_dir / f"{name}.so").resolve())
+    lines.append(f"build {output_so}: {link_rule} " + " ".join(objects))
+    lines.append(f"default {output_so}")
     lines.append("")
 
     return "\n".join(lines)