refactor: Builder System #120
Conversation
WalkthroughA comprehensive refactoring of the build system replaces cache and builder patterns with a centralized BuilderRegistry, introduces a new RunnableMetadata model for build artifacts, removes the CUDA builder, adds a TorchBuilder replacement, updates all builders to use cache-path semantics, and adds utility functions for package naming and file I/O. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @Ubospica, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a comprehensive refactoring of the builder system, which is central to how solution implementations are compiled and executed. The changes aim to enhance modularity, improve extensibility, and optimize performance through intelligent caching. Key updates include the introduction of a dedicated Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a major and well-executed refactoring of the builder system. The new design centralizes caching in a BuilderRegistry, simplifies the Builder interface, and improves the structure of concrete builder implementations. The introduction of RunnableMetadata and a unified Runnable class with call_dps support is a great improvement.
I've found a few critical issues that will cause runtime errors: a typo in the Runnable.cleanup method and a missing value in the BuildType literal that will break TorchBuilder. I also noticed that several tests seem to be broken due to the refactoring, as they call a non-existent build_with_cache method. Additionally, I have a couple of suggestions to improve code style and maintainability in TritonBuilder and TVMFFIBuilder.
Overall, this is a high-quality refactoring. Once the identified issues are addressed, the builder system will be much more robust and maintainable.
| from flashinfer_bench.data import Definition | ||
| from flashinfer_bench.utils import dtype_str_to_torch_dtype | ||
|
|
||
| BuildType = Literal["cuda", "tvm_ffi", "python", "triton"] |
There was a problem hiding this comment.
The BuildType literal is defined as Literal["cuda", "tvm_ffi", "python", "triton"]. However, the new TorchBuilder sets the build_type to "torch". This will cause a pydantic validation error at runtime. The BuildType should be updated to include "torch". Since CUDABuilder has been removed and replaced by TorchBuilder, "cuda" should probably be replaced with "torch".
| BuildType = Literal["cuda", "tvm_ffi", "python", "triton"] | |
| BuildType = Literal["torch", "tvm_ffi", "python", "triton"] |
| if self._closer: | ||
| try: | ||
| self._closer() | ||
| finally: | ||
| self._closer = None |
There was a problem hiding this comment.
There's a typo in the cleanup method. The __init__ method sets the cleaner function to self._cleaner, but this method attempts to access self._closer. This will result in an AttributeError when cleanup is called.
| if self._closer: | |
| try: | |
| self._closer() | |
| finally: | |
| self._closer = None | |
| if self._cleaner: | |
| try: | |
| self._cleaner() | |
| finally: | |
| self._cleaner = None |
| r1 = b.build_with_cache(d, s) | ||
| r2 = b.build_with_cache(d, s) |
There was a problem hiding this comment.
This test, along with others in test_python_builder.py, test_triton_builder.py, and test_tvm_ffi_builder.py, has been updated to call build_with_cache. However, this method does not exist on the Builder class after the refactoring. The caching logic has been moved to BuilderRegistry, and the build method on Builder no longer performs caching.
These tests appear to be broken and will fail. They need to be updated to reflect the new design, likely by testing the BuilderRegistry for caching behavior or by adapting the builder tests to call the build method directly without expecting caching at the builder level.
Additionally, DummyBuilder in this file implements _build but the Builder abstract base class now requires build to be implemented, so DummyBuilder is no longer a valid subclass of Builder.
| def build(self, definition: Definition, solution: Solution) -> Runnable: | ||
| """Build a solution into a runnable implementation. | ||
| This method compiles/loads the solution's source code and returns a Runnable | ||
| object that can be executed with the interface specified by the definition. | ||
| Parameters | ||
| ---------- | ||
| definition : Definition | ||
| The problem definition that specifies the expected interface. | ||
| solution : Solution | ||
| The solution implementation to build. | ||
| Returns | ||
| ------- | ||
| Runnable | ||
| An executable wrapper around the built implementation. | ||
| Raises | ||
| ------ | ||
| BuildError | ||
| If the build fails for any reason (compilation errors, missing dependencies, etc.). | ||
| """ | ||
| ... |
There was a problem hiding this comment.
The build method is now abstract, but some test implementations like DummyBuilder in tests/compile/test_builder.py still implement _build, which will cause an error. Please ensure all concrete builder implementations, including those in tests, are updated to implement the public build method.
| super().__init__() | ||
| self._py_builder = py_builder | ||
| def __init__(self) -> None: | ||
| Builder.__init__(self, self._KEY_PREFIX, self._BUILD_DIR_NAME) |
There was a problem hiding this comment.
The __init__ method directly calls Builder.__init__, bypassing the __init__ of its direct parent, PythonBuilder. While this works currently because PythonBuilder.__init__ only calls its superclass's __init__, it's a fragile design. If PythonBuilder.__init__ is modified in the future to include additional setup logic, TritonBuilder would not inherit it. It's more robust to use super().__init__() to respect the method resolution order and make the inheritance chain clearer.
| Builder.__init__(self, self._KEY_PREFIX, self._BUILD_DIR_NAME) | |
| super().__init__() |
| import tvm_ffi | ||
| from tvm_ffi.utils import FileLock |
There was a problem hiding this comment.
The imports for tvm_ffi and tvm_ffi.utils.FileLock are inside the build method. It's generally better practice to place imports at the top of the file for clarity and to avoid repeated import overhead. Since the availability of tvm_ffi is already checked by is_available() before this builder is used, these imports can be safely moved to the top level of the file.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (12)
flashinfer_bench/compile/__init__.py (1)
16-27: Exports look correct; consider sorting__all__alphabetically.The new
RunnableMetadataexport is appropriate. Static analysis flags that__all__is not sorted (RUF022). If the project enforces isort-style sorting, apply this diff:__all__ = [ "Builder", "BuildError", "BuilderRegistry", + "get_builder_registry", "Runnable", "RunnableMetadata", - "get_builder_registry", ]flashinfer_bench/compile/builders/__init__.py (1)
1-8: LGTM!TorchBuildercorrectly replacesCUDABuilder.The module docstring is helpful and the export list is updated correctly. If the project enforces sorted
__all__(RUF022), consider alphabetical ordering:-__all__ = ["TorchBuilder", "PythonBuilder", "TritonBuilder", "TVMFFIBuilder"] +__all__ = ["PythonBuilder", "TorchBuilder", "TritonBuilder", "TVMFFIBuilder"]flashinfer_bench/compile/utils.py (2)
39-54: Security checks use assertions - consider using explicit exceptions.Assertions can be disabled with
python -O, which would bypass these security checks. TheSourceFilemodel already validates paths at creation time (per the relevant snippet fromflashinfer_bench/data/solution.py), so these are truly defensive. However, for defense-in-depth, explicit exceptions are safer:- assert not src_path_obj.is_absolute(), f"Absolute path detected: {src.path}" - assert ".." not in src_path_obj.parts, f"Path traversal detected: {src.path}" + if src_path_obj.is_absolute(): + raise ValueError(f"Absolute path detected: {src.path}") + if ".." in src_path_obj.parts: + raise ValueError(f"Path traversal detected: {src.path}")
28-31: Minor: Docstring says "absolute paths" but returned paths may not be resolved.The docstring states the function returns "List of absolute paths" but the code appends
path / src.pathwithout calling.resolve(). Ifpathis relative, the returned paths will also be relative.- paths.append(src_path) + paths.append(src_path.resolve())Alternatively, update the docstring to say "paths relative to the provided root" if that's the intended behavior.
Also applies to: 53-53
flashinfer_bench/compile/builders/torch_builder.py (1)
96-131: Robust error handling; consider using spread operator for list construction.The broad
Exceptioncatch (line 125) is appropriate here since dependency discovery is best-effort. For the Windows linker flags (line 123), using the spread operator is cleaner per RUF005:- ldflags = [f"/LIBPATH:{lib_path}"] + lib_names + ldflags = [f"/LIBPATH:{lib_path}", *lib_names]flashinfer_bench/compile/builder.py (1)
7-7: Unused importDict.
Dictis imported but not used in this file.-from typing import Dict, Tuple +from typing import Tupleflashinfer_bench/compile/builders/python_builder.py (2)
6-6: Unused importos.
osis imported but not used in this file.import importlib -import os import shutil
41-58:get_keyand_get_build_pathduplicate base class logic.These methods replicate functionality already provided by
get_package_name_and_build_pathin the base class. Consider removing them to reduce code duplication, unless there's a specific reason to keep them separate.flashinfer_bench/compile/registry.py (1)
122-124: Variablehashshadows builtin.Using
hashas a variable name shadows the Python builtin. Consider renaming tosol_hashorsolution_hash.- hash = sol.hash() - if hash in self._cache: - return self._cache[hash] + sol_hash = sol.hash() + if sol_hash in self._cache: + return self._cache[sol_hash]And update the cache assignment on line 130:
- self._cache[hash] = runnable + self._cache[sol_hash] = runnableflashinfer_bench/compile/builders/triton_builder.py (2)
28-29: Direct call toBuilder.__init__bypassesPythonBuilder.__init__.Calling
Builder.__init__(self, ...)directly instead ofsuper().__init__()is fragile. IfPythonBuilder.__init__ever adds initialization logic beyond calling its parent,TritonBuilderwill miss it.Since
PythonBuilder.__init__currently just callssuper().__init__(), this works, but it's a maintenance hazard.Consider overriding the class variables and letting the parent chain handle initialization:
class TritonBuilder(PythonBuilder): ... _KEY_PREFIX: ClassVar[str] = "fib_triton_" _BUILD_DIR_NAME: ClassVar[str] = "triton" def __init__(self) -> None: - Builder.__init__(self, self._KEY_PREFIX, self._BUILD_DIR_NAME) + super().__init__()But this requires
PythonBuilder.__init__to useself._KEY_PREFIXandself._BUILD_DIR_NAMEinstead of hardcoded class variables:# In PythonBuilder.__init__: def __init__(self) -> None: super().__init__(self._KEY_PREFIX, self._BUILD_DIR_NAME)
68-69: Mutatingmetadata.build_typeafter construction.While this works because
RunnableMetadataisn't frozen, mutating the returnedRunnable's metadata is somewhat surprising. Consider passing the correctbuild_typeduring construction instead.An alternative approach would be to override the metadata construction in the build method rather than mutating it post-hoc. However, since this would require more significant refactoring of
PythonBuilder.build(), the current approach is acceptable.flashinfer_bench/compile/builders/tvm_ffi_builder.py (1)
119-122: Consider using explicit checks instead of assertions for path validation.Assertions can be disabled with
python -O. While the comment notes this is defensive (relying on upstream validation), if this code could run in optimized mode, these security checks would be bypassed.# Defensive assertion: the path in the solution should be validated by the Solution # model validator, but we add this defensive assertion to be safe. src_path_obj = Path(src.path) - assert not src_path_obj.is_absolute(), f"Absolute path detected: {src.path}" - assert ".." not in src_path_obj.parts, f"Path traversal detected: {src.path}" + if src_path_obj.is_absolute(): + raise BuildError(f"Absolute path detected: {src.path}") + if ".." in src_path_obj.parts: + raise BuildError(f"Path traversal detected: {src.path}")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (22)
examples/win_at_p.py(0 hunks)flashinfer_bench/apply/key.py(1 hunks)flashinfer_bench/compile/__init__.py(1 hunks)flashinfer_bench/compile/builder.py(1 hunks)flashinfer_bench/compile/builders/__init__.py(1 hunks)flashinfer_bench/compile/builders/cuda_builder.py(0 hunks)flashinfer_bench/compile/builders/python_builder.py(1 hunks)flashinfer_bench/compile/builders/torch_builder.py(1 hunks)flashinfer_bench/compile/builders/triton_builder.py(1 hunks)flashinfer_bench/compile/builders/tvm_ffi_builder.py(8 hunks)flashinfer_bench/compile/registry.py(2 hunks)flashinfer_bench/compile/runnable.py(2 hunks)flashinfer_bench/compile/utils.py(1 hunks)flashinfer_bench/data/solution.py(5 hunks)flashinfer_bench/data/trace_set.py(0 hunks)tests/apply/test_runtime.py(2 hunks)tests/compile/test_builder.py(2 hunks)tests/compile/test_python_builder.py(2 hunks)tests/compile/test_runnable.py(2 hunks)tests/compile/test_triton_builder.py(3 hunks)tests/compile/test_tvm_ffi_builder.py(8 hunks)tests/data/test_solution.py(0 hunks)
💤 Files with no reviewable changes (4)
- flashinfer_bench/data/trace_set.py
- examples/win_at_p.py
- flashinfer_bench/compile/builders/cuda_builder.py
- tests/data/test_solution.py
🧰 Additional context used
🧬 Code graph analysis (8)
flashinfer_bench/compile/__init__.py (3)
flashinfer_bench/compile/builder.py (2)
Builder(20-93)BuildError(16-17)flashinfer_bench/compile/registry.py (2)
BuilderRegistry(25-167)get_builder_registry(170-180)flashinfer_bench/compile/runnable.py (2)
Runnable(33-172)RunnableMetadata(16-30)
flashinfer_bench/compile/utils.py (1)
flashinfer_bench/data/solution.py (2)
SourceFile(30-61)hash(186-213)
flashinfer_bench/compile/builders/python_builder.py (6)
flashinfer_bench/compile/builder.py (2)
build(54-77)get_package_name_and_build_path(79-93)flashinfer_bench/compile/runnable.py (2)
Runnable(33-172)RunnableMetadata(16-30)flashinfer_bench/compile/utils.py (2)
create_package_name(58-91)write_sources_to_path(12-55)flashinfer_bench/data/definition.py (1)
Definition(95-363)flashinfer_bench/data/solution.py (4)
Solution(100-213)SupportedLanguages(13-27)get_entry_path(145-156)get_entry_symbol(158-170)flashinfer_bench/env.py (1)
get_fib_cache_path(46-57)
flashinfer_bench/compile/builders/triton_builder.py (5)
flashinfer_bench/compile/builder.py (3)
Builder(20-93)can_build(38-51)build(54-77)flashinfer_bench/data/solution.py (2)
Solution(100-213)SupportedLanguages(13-27)flashinfer_bench/compile/builders/python_builder.py (2)
can_build(37-39)build(97-162)flashinfer_bench/compile/builders/torch_builder.py (3)
is_available(135-147)can_build(149-151)build(241-311)flashinfer_bench/compile/registry.py (1)
build(97-132)
flashinfer_bench/compile/runnable.py (5)
flashinfer_bench/data/definition.py (2)
get_var_values(239-279)get_output_shapes(343-363)flashinfer_bench/utils.py (1)
dtype_str_to_torch_dtype(39-45)flashinfer_bench/compile/builders/python_builder.py (1)
cleaner(79-93)flashinfer_bench/compile/builders/torch_builder.py (1)
cleaner(236-237)flashinfer_bench/compile/builders/tvm_ffi_builder.py (1)
cleaner(201-202)
flashinfer_bench/compile/builders/__init__.py (4)
flashinfer_bench/compile/builders/python_builder.py (1)
PythonBuilder(19-162)flashinfer_bench/compile/builders/torch_builder.py (1)
TorchBuilder(31-311)flashinfer_bench/compile/builders/triton_builder.py (1)
TritonBuilder(14-70)flashinfer_bench/compile/builders/tvm_ffi_builder.py (1)
TVMFFIBuilder(23-296)
tests/compile/test_builder.py (2)
flashinfer_bench/compile/builders/python_builder.py (2)
get_key(41-43)cleaner(79-93)flashinfer_bench/compile/runnable.py (1)
Runnable(33-172)
tests/apply/test_runtime.py (5)
flashinfer_bench/compile/builders/python_builder.py (2)
PythonBuilder(19-162)build(97-162)flashinfer_bench/compile/builders/torch_builder.py (1)
build(241-311)flashinfer_bench/compile/registry.py (1)
build(97-132)flashinfer_bench/data/solution.py (1)
Solution(100-213)flashinfer_bench/compile/runnable.py (1)
Runnable(33-172)
🪛 GitHub Actions: .github/workflows/linting.yaml
flashinfer_bench/compile/builders/tvm_ffi_builder.py
[error] 59-59: ruff: F401 'tvm_ffi' imported but unused. Remove unused import.
🪛 Ruff (0.14.6)
flashinfer_bench/compile/__init__.py
20-27: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
flashinfer_bench/data/solution.py
202-202: Probable use of insecure hash functions in hashlib: sha1
(S324)
flashinfer_bench/compile/builders/python_builder.py
85-86: try-except-pass detected, consider logging the exception
(S110)
85-85: Do not catch blind exception: Exception
(BLE001)
125-125: Avoid specifying long messages outside the exception class
(TRY003)
141-141: Avoid specifying long messages outside the exception class
(TRY003)
147-149: Avoid specifying long messages outside the exception class
(TRY003)
153-153: Avoid specifying long messages outside the exception class
(TRY003)
flashinfer_bench/compile/registry.py
58-58: Avoid specifying long messages outside the exception class
(TRY003)
72-73: try-except-pass detected, consider logging the exception
(S110)
72-72: Do not catch blind exception: Exception
(BLE001)
132-132: Avoid specifying long messages outside the exception class
(TRY003)
flashinfer_bench/compile/builders/torch_builder.py
123-123: Consider [f"/LIBPATH:{lib_path}", *lib_names] instead of concatenation
Replace with [f"/LIBPATH:{lib_path}", *lib_names]
(RUF005)
125-125: Do not catch blind exception: Exception
(BLE001)
211-214: Avoid specifying long messages outside the exception class
(TRY003)
269-272: Avoid specifying long messages outside the exception class
(TRY003)
295-295: Avoid specifying long messages outside the exception class
(TRY003)
300-300: Avoid specifying long messages outside the exception class
(TRY003)
flashinfer_bench/compile/runnable.py
128-131: Avoid specifying long messages outside the exception class
(TRY003)
flashinfer_bench/compile/builders/__init__.py
8-8: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
tests/compile/test_builder.py
27-27: Unused method argument: definition
(ARG002)
27-27: Unused method argument: solution
(ARG002)
flashinfer_bench/compile/builders/tvm_ffi_builder.py
267-269: Avoid specifying long messages outside the exception class
(TRY003)
293-293: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Run unit tests on ubuntu-latest and Python 3.10
- GitHub Check: Run unit tests on ubuntu-latest and Python 3.12
- GitHub Check: Run unit tests on ubuntu-latest and Python 3.11
- GitHub Check: Run unit tests on ubuntu-latest and Python 3.13
🔇 Additional comments (36)
flashinfer_bench/apply/key.py (1)
15-15: Clean type annotation refactoring.Flattening the nested Union to
Union[int, float, bool]is more idiomatic and improves readability without affecting runtime behavior.flashinfer_bench/data/solution.py (2)
24-27: LGTM! New CPP language variant properly added.The new
CPPenum value follows the existing pattern and documentation style.
145-170: LGTM! Clean helper methods for entry point parsing.Both
get_entry_path()andget_entry_symbol()properly extract components from the validated entry point format. Usingsplit("::")[-1]for the symbol is a good defensive choice.tests/compile/test_python_builder.py (2)
41-41: LGTM! Correctly updated to use cache-aware build API.The test properly sets up an isolated cache directory and uses the new
build_with_cachemethod.
80-80: LGTM! Consistent API migration.Same pattern applied correctly with proper cache cleanup via
b.clear_cache().tests/compile/test_builder.py (3)
21-22: LGTM! Method renamed to match updated builder interface.The
get_keynaming is consistent with the productionPythonBuilder.get_keymethod shown in the relevant code snippets.
48-50: LGTM! Cache behavior correctly tested.The test verifies that
build_with_cachereturns the same instance on repeated calls (r1 is r2), confirming cache hit behavior.
27-30: Verifymetadataparameter type compatibility.The
Runnableconstructor expectsmetadata: RunnableMetadata, but the test passes a plain dict{"dummy": True}. Confirm whetherRunnableMetadataaccepts dict input (e.g., via TypedDict or dataclass supporting dict initialization). IfRunnableMetadatais aTypedDict, this usage is valid; if it's a stricter type, adjust the test accordingly.tests/compile/test_triton_builder.py (3)
57-57: LGTM! Import guard test correctly uses cache-aware API.The test verifies that
BuildErroris raised when Triton is unavailable.
77-77: LGTM! Minimum test updated consistently.Proper cache directory setup and availability state reset via
monkeypatch.setattr.
141-141: LGTM! Vector add test with correct API usage.The Triton kernel test properly exercises the cached build path and verifies execution correctness.
tests/compile/test_tvm_ffi_builder.py (8)
123-123: LGTM! CPU build test migrated to cache-aware API.Test correctly uses the
isolated_cachefixture and verifies kernel execution.
170-170: LGTM! CUDA GPU test updated consistently.
254-262: LGTM! Builder-level caching test properly validates cache behavior.The test measures build time vs cache load time, demonstrating the caching benefit.
303-311: LGTM! Cross-builder caching test correctly validates shared cache.Tests that different
TVMFFIBuilderinstances can share cached artifacts.
353-353: LGTM! call_dest test updated.
404-404: LGTM! Error handling test for invalid entry point.
436-436: LGTM! Error handling test for invalid sources.
469-469: LGTM! Subdirectory source handling test updated.flashinfer_bench/compile/__init__.py (1)
1-14: LGTM! Clear and helpful docstring.The module docstring provides a good overview of the package components and a concise workflow example.
tests/apply/test_runtime.py (1)
246-272: LGTM! Correctly updated to patch the publicbuildmethod.The change from
_buildtobuildaligns with the refactored builder API. The patch/restore logic is correct.Note: This test is currently skipped (line 163). Ensure the underlying issue is tracked for resolution.
tests/compile/test_runnable.py (2)
5-5: LGTM! Import updated correctly.
17-27: Test correctly uses the new Runnable API, but there's a potential bug in theRunnable.cleanup()method.The test construction and cleanup calls are correct. However, the
Runnableclass may have an inconsistency where the constructor assigns toself._cleanerbut thecleanup()method referencesself._closer. This mismatch would cause the test to fail at runtime with anAttributeErrorsince_closerwould be undefined.flashinfer_bench/compile/utils.py (1)
58-91: LGTM! Package name generation is sound.The normalization handles edge cases well (digit prefix, special characters). Using the first 6 characters of the SHA1 hash provides reasonable uniqueness for cache keys.
flashinfer_bench/compile/builders/torch_builder.py (2)
50-68: LGTM! Clean dependency discovery pattern.The lazy discovery of CUDA dependencies at init time with graceful fallback is appropriate.
284-311: Build logic is correct and well-structured.The compilation flow using
torch.utils.cpp_extension.load()is appropriate, with proper error handling and metadata construction. The cleaner pattern ensures resources can be released.flashinfer_bench/compile/builder.py (1)
20-93: Well-structured abstract base class.The
Builderabstraction is clean and provides a good foundation for concrete builders. Theget_package_name_and_build_pathhelper centralizes path computation, reducing duplication across subclasses.flashinfer_bench/compile/builders/python_builder.py (2)
79-95: Cleaner implementation is sound.The cleaner correctly handles module unloading, sys.path cleanup, and directory removal. The broad
except Exception: passis acceptable here since cleanup failures shouldn't propagate and prevent further cleanup steps.
97-162: Build method is well-structured with proper error handling.The build flow correctly validates the entry file, writes sources, imports the module, and constructs the Runnable with appropriate metadata. Error handling properly invokes the cleaner before raising
BuildError.flashinfer_bench/compile/registry.py (3)
62-74: Clear method correctly handles cleanup failures.The broad exception catch is appropriate here to ensure all runnables are processed even if some cleanup operations fail.
134-167:build_referencecreates a well-formed pseudo-solution.The method correctly constructs a Solution object with appropriate defaults for building reference implementations.
89-94: Missingis_available()in Builder interface.
builder_type.is_available()is called on line 92, butis_available()is not defined as an abstract method in theBuilderbase class. If any builder subclass doesn't implement this method, this will raise anAttributeErrorat runtime.Consider adding
is_available()as an abstract static method inBuilder, or add a default implementation that returnsTrue:# In builder.py + @staticmethod + def is_available() -> bool: + """Check if this builder is available in the current environment. + + Returns True by default. Subclasses should override if they have + optional dependencies. + """ + return Trueflashinfer_bench/compile/builders/triton_builder.py (1)
31-48:is_available()andcan_build()are correctly implemented.The availability check properly handles import failures, and
can_buildcorrectly filters by language type.flashinfer_bench/compile/runnable.py (2)
16-31:RunnableMetadatamodel is well-defined.The Pydantic model provides good structure for metadata with appropriate fields and documentation.
93-159:call_dpsimplementation is solid.The DPS wrapper correctly:
- Validates the presence of the definition in metadata
- Derives output shapes from input shapes
- Handles device placement
- Normalizes return values (None/single/tuple)
flashinfer_bench/compile/builders/tvm_ffi_builder.py (1)
206-296: LGTM!The build method is well-implemented with proper multi-process safety through the FileLock and double-check pattern after acquiring the lock. The caching logic correctly handles concurrent builds, and the Runnable construction with metadata and cleaner follows the new builder pattern.
| package_name, build_path = self.get_package_name_and_build_path(solution) | ||
| module_name = package_name + "." + ".".join(Path(entry_file).with_suffix("").parts) | ||
|
|
||
| build_path = self._get_build_path(package_name) |
There was a problem hiding this comment.
Redundant build_path computation.
get_package_name_and_build_path already returns build_path, but it's immediately overwritten by _get_build_path. This is wasteful and confusing. Either use the base class method entirely, or remove the call to get_package_name_and_build_path and just use get_key + _get_build_path.
- package_name, build_path = self.get_package_name_and_build_path(solution)
+ package_name, build_path = self.get_package_name_and_build_path(solution)
module_name = package_name + "." + ".".join(Path(entry_file).with_suffix("").parts)
- build_path = self._get_build_path(package_name)Alternatively, if _get_build_path logic differs intentionally, use:
- package_name, build_path = self.get_package_name_and_build_path(solution)
+ package_name = self.get_key(solution)
module_name = package_name + "." + ".".join(Path(entry_file).with_suffix("").parts)
build_path = self._get_build_path(package_name)Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In flashinfer_bench/compile/builders/python_builder.py around lines 127 to 130,
there's a redundant call to get_package_name_and_build_path whose returned
build_path is immediately overwritten by _get_build_path; remove the unnecessary
value and either (A) call get_package_name_and_build_path to obtain both
package_name and build_path and stop calling _get_build_path, or (B) if you only
need package_name here, replace get_package_name_and_build_path with get_key (or
another method that returns only the package_name) and then call
_get_build_path(package_name); update the surrounding code to use the single,
consistent source of truth for build_path.
| extra_include_paths, extra_ldflags = self._get_dependency_flags(solution) | ||
| # Add build directory to include paths | ||
| extra_include_paths.append(build_dir) |
There was a problem hiding this comment.
Type mismatch: build_dir is Path but extra_include_paths is List[str].
extra_include_paths collects strings, but build_dir is a Path object. This may cause issues with torch.utils.cpp_extension.load() depending on how it handles the list. Convert to string for consistency:
# Add build directory to include paths
- extra_include_paths.append(build_dir)
+ extra_include_paths.append(str(build_dir))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| extra_include_paths, extra_ldflags = self._get_dependency_flags(solution) | |
| # Add build directory to include paths | |
| extra_include_paths.append(build_dir) | |
| extra_include_paths, extra_ldflags = self._get_dependency_flags(solution) | |
| # Add build directory to include paths | |
| extra_include_paths.append(str(build_dir)) |
🤖 Prompt for AI Agents
In flashinfer_bench/compile/builders/torch_builder.py around lines 280 to 282,
extra_include_paths is a list of strings but build_dir is a Path object; convert
build_dir to a string before appending so the list is homogeneous (e.g., append
str(build_dir)) to avoid type errors when passing extra_include_paths to
torch.utils.cpp_extension.load().
| @staticmethod | ||
| def is_available() -> bool: | ||
| """Check if TVM-FFI is available in the current environment.""" | ||
| try: | ||
| import tvm_ffi | ||
| except ImportError: | ||
| return False | ||
| return True |
There was a problem hiding this comment.
Fix unused import linter error (F401) breaking CI.
The pipeline is failing because ruff sees the import tvm_ffi as unused. Add a noqa comment to suppress this false positive since the import is intentionally used only for availability checking.
@staticmethod
def is_available() -> bool:
"""Check if TVM-FFI is available in the current environment."""
try:
- import tvm_ffi
+ import tvm_ffi # noqa: F401
except ImportError:
return False
return True🧰 Tools
🪛 GitHub Actions: .github/workflows/linting.yaml
[error] 59-59: ruff: F401 'tvm_ffi' imported but unused. Remove unused import.
🤖 Prompt for AI Agents
In flashinfer_bench/compile/builders/tvm_ffi_builder.py around lines 56 to 63,
the try/except import of tvm_ffi is flagged by the linter as an unused import
(F401); modify the import line to suppress this false positive by adding a noqa
comment (e.g., "import tvm_ffi # noqa: F401") so the availability check remains
intact and the linter error is silenced.
| BuildType = Literal["cuda", "tvm_ffi", "python", "triton"] | ||
| """The type of build that produced this runnable. Each builder has a unique build type.""" |
There was a problem hiding this comment.
BuildType is missing "torch".
TorchBuilder sets build_type="torch" (per torch_builder.py line 299), but BuildType only includes "cuda", "tvm_ffi", "python", "triton". This will cause Pydantic validation errors when building CUDA solutions with TorchBuilder.
-BuildType = Literal["cuda", "tvm_ffi", "python", "triton"]
+BuildType = Literal["python", "triton", "torch", "tvm_ffi"]Also update the docstring on line 24 which mentions "torch" but it's not currently in the type.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| BuildType = Literal["cuda", "tvm_ffi", "python", "triton"] | |
| """The type of build that produced this runnable. Each builder has a unique build type.""" | |
| BuildType = Literal["python", "triton", "torch", "tvm_ffi"] | |
| """The type of build that produced this runnable. Each builder has a unique build type.""" |
🤖 Prompt for AI Agents
In flashinfer_bench/compile/runnable.py around lines 12-13 (and update the
docstring at line 24), the BuildType Literal is missing the "torch" variant
which causes Pydantic validation failures because TorchBuilder uses
build_type="torch"; add "torch" to the Literal union for BuildType and update
the docstring on line 24 to include "torch" in the listed build types so the
type and documentation match.
| if self._closer: | ||
| try: | ||
| self._closer() | ||
| finally: | ||
| self._closer = None |
There was a problem hiding this comment.
Critical bug: _closer should be _cleaner.
The cleanup() method references self._closer, but the field is defined as self._cleaner (lines 46, 68). This will raise an AttributeError when cleanup() is called.
def cleanup(self) -> None:
"""Clean up build artifacts and release resources.
This method calls the cleaner function if one was provided during construction.
It is idempotent: calling it multiple times is safe and has no additional effect
after the first call.
"""
- if self._closer:
+ if self._cleaner:
try:
- self._closer()
+ self._cleaner()
finally:
- self._closer = None
+ self._cleaner = None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if self._closer: | |
| try: | |
| self._closer() | |
| finally: | |
| self._closer = None | |
| if self._cleaner: | |
| try: | |
| self._cleaner() | |
| finally: | |
| self._cleaner = None |
🤖 Prompt for AI Agents
In flashinfer_bench/compile/runnable.py around lines 168 to 172, the cleanup()
method mistakenly references self._closer (which does not exist) instead of the
defined self._cleaner; update the block to call self._cleaner() inside the try
and then set self._cleaner = None in the finally so the cleanup uses the correct
attribute and clears it afterward.
| def hash(self) -> str: | ||
| """Compute a deterministic hash of the solution content. | ||
| The hash is computed from all fields that affect the solution's behavior: | ||
| name, definition, language, entry point, dependencies, and all source file | ||
| paths and contents. This ensures that any meaningful change to the solution | ||
| results in a different hash. | ||
| The hash is used for caching build artifacts - solutions with the same hash | ||
| can reuse the same cached build result. | ||
| Returns | ||
| ------- | ||
| bool | ||
| True if the solution requires building (has build commands or uses CUDA), | ||
| False otherwise. | ||
| str | ||
| A SHA1 hash (40 hex characters) uniquely identifying this solution's content. | ||
| """ | ||
| return self.spec.language == SupportedLanguages.CUDA | ||
| h = hashlib.sha1() | ||
| for s in ( | ||
| self.name, | ||
| self.definition, | ||
| self.spec.language, | ||
| self.spec.entry_point, | ||
| *self.spec.dependencies, | ||
| *(part for src in self.sources for part in (src.path, src.content)), | ||
| ): | ||
| h.update(s.encode()) | ||
|
|
||
| return h.hexdigest() |
There was a problem hiding this comment.
Potential TypeError if dependencies is None.
The dependencies field is typed as Optional[List[NonEmptyString]], meaning it can be explicitly set to None. If that happens, line 208 will raise TypeError: 'NoneType' object is not iterable when unpacking with *self.spec.dependencies.
h = hashlib.sha1()
for s in (
self.name,
self.definition,
self.spec.language,
self.spec.entry_point,
- *self.spec.dependencies,
+ *(self.spec.dependencies or []),
*(part for src in self.sources for part in (src.path, src.content)),
):
h.update(s.encode())Regarding the static analysis hint about SHA1 (S324): SHA1 is appropriate here since it's used for content-based cache keying, not cryptographic security. The concern is a false positive in this context.
🧰 Tools
🪛 Ruff (0.14.6)
202-202: Probable use of insecure hash functions in hashlib: sha1
(S324)
🤖 Prompt for AI Agents
In flashinfer_bench/data/solution.py around lines 186 to 213, the hash() method
unpacks self.spec.dependencies with *self.spec.dependencies which will raise
TypeError if dependencies is None; change to iterate over self.spec.dependencies
or an empty iterable (e.g. self.spec.dependencies or []) when building the tuple
of fields so None is treated as no dependencies, and ensure each item appended
is a str before calling .encode() (cast or skip None) so h.update(...) always
receives bytes.
1 participant
Signed-off-by: Ubospica ubospica@gmail.com
Summary by CodeRabbit
New Features
Refactoring
Removals
✏️ Tip: You can customize this high-level summary in your review settings.