[To merge AFTER flashinfer-ci changes updated] Reduce test time by moving compilation off-line #2089
Conversation
WalkthroughAdds non-dry-run runtime initialization to install and verify prebuilt kernel artifacts and local Python sources, and simplifies the MAX_JOBS calculation for JIT cache wheel builds by using a constant divisor for memory-based job estimation. Changes
Sequence Diagram(s)sequenceDiagram
participant Script as task_test_blackwell_kernels.sh
participant Env as Environment
participant FS as Filesystem (dist/)
participant Pip as pip
participant Python as python
rect rgb(240,248,255)
Note over Script: Start (only if DRY_RUN unset)
Script->>Env: Check DRY_RUN
alt DRY_RUN not set
Script->>Env: Echo CUDA_VERSION
Script->>Script: Compute JIT_ARCH_EFFECTIVE (special-case 12.0/cu129)
Script->>Script: Set DIST_CUBIN_DIR = dist/${CUDA_VERSION}/cubin
Script->>Script: Set DIST_JIT_CACHE_DIR = dist/${CUDA_VERSION}/jit-cache
else DRY_RUN set
Script-->>Env: Skip install steps
end
end
rect rgb(230,255,240)
Note over Script,FS: Install prebuilt artifacts (error if missing)
Script->>FS: Check DIST_CUBIN_DIR for wheels
alt cubin wheels found
Script->>Pip: pip install <cubin wheel>
Pip-->>Script: success
else missing
Script-->>Env: exit with error ("missing flashinfer-cubin artifact")
end
Script->>FS: Check DIST_JIT_CACHE_DIR for wheels
alt jit-cache wheels found
Script->>Pip: pip install <jit-cache wheel>
Pip-->>Script: success
else missing
Script-->>Env: exit with error ("missing flashinfer-jit-cache artifact")
end
end
rect rgb(255,250,230)
Note over Script,Python: Install local package & verify
Script->>Pip: pip install -e . -v --no-deps
Pip-->>Script: installed
Script->>Python: (cd /tmp && python -m flashinfer show-config)
Python-->>Script: config output / verification
end
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 📝 Customizable high-level summaries are now available in beta!You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
Example instruction:
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later. 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 |
kahyunnam
force-pushed
the
knam/unit-testing-move-compilation-offline
branch
from
c9c6768 to
375ca18
Compare
|
/bot run |
kahyunnam
requested review from
bkryu,
nvmbreughe,
yongwww and
yzh119
as code owners
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
scripts/task_test_blackwell_kernels.sh (2)
41-50: Inconsistent verbosity flags in sequential pip installations.Lines 43 and 45 use
-q(quiet) for kernel installations, while line 49 uses-v(verbose) for local source installation. This inconsistency makes it unclear whether the verbosity change is intentional and may make output harder to parse in CI logs.Standardize the verbosity flags across all pip installations in this initialization block:
# Install precompiled kernels echo "Installing flashinfer-cubin from PyPI/index..." - pip install -q flashinfer-cubin + pip install -q flashinfer-cubin echo "Installing flashinfer-jit-cache for ${CUDA_STREAM} from https://flashinfer.ai/whl/${CUDA_STREAM} ..." - pip install -q --extra-index-url "https://flashinfer.ai/whl/${CUDA_STREAM}" flashinfer-jit-cache + pip install -q --extra-index-url "https://flashinfer.ai/whl/${CUDA_STREAM}" flashinfer-jit-cache echo "" # Install local python sources - pip install -e . -v --no-deps + pip install -e . -q --no-depsAlternatively, if verbose output is intentional for debugging local installs, add a comment explaining the choice.
41-50: Verify that the custom PyPI index URL forflashinfer-jit-cacheis reliable.The script hardcodes the index URL
https://flashinfer.ai/whl/${CUDA_STREAM}and expects it to always be available and contain theflashinfer-jit-cachepackage for the detected CUDA stream. If this URL becomes unavailable or if a CUDA stream version is not published, the pip install will fail and halt all subsequent tests.Add error handling and diagnostics to surface issues clearly:
echo "Installing flashinfer-jit-cache for ${CUDA_STREAM} from https://flashinfer.ai/whl/${CUDA_STREAM} ..." - pip install -q --extra-index-url "https://flashinfer.ai/whl/${CUDA_STREAM}" flashinfer-jit-cache + if ! pip install -q --extra-index-url "https://flashinfer.ai/whl/${CUDA_STREAM}" flashinfer-jit-cache; then + echo "❌ ERROR: Failed to install flashinfer-jit-cache for CUDA stream ${CUDA_STREAM}" + echo " Index URL: https://flashinfer.ai/whl/${CUDA_STREAM}" + exit 1 + fiCan you confirm that the custom index URL is stable and that all supported CUDA streams (cu128, cu129, cu130) are consistently published with the corresponding flashinfer-jit-cache package?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
scripts/task_test_blackwell_kernels.sh(1 hunks)
🔇 Additional comments (1)
scripts/task_test_blackwell_kernels.sh (1)
52-55: Verify thatpython -m flashinfer show-configis an appropriate verification step.The verification runs
python -m flashinfer show-configto confirm successful installation. However, this assumes:
- The
show-configsubcommand exists in the flashinfer module- The command is idempotent and doesn't modify the environment
- The command completes quickly without external dependencies
If this command fails (e.g., due to missing dependencies, invalid environment, or a transient issue), the entire test run is aborted before any tests can run, which may be overly strict for a verification step.
Can you confirm:
- That
python -m flashinfer show-configis a lightweight, read-only command that verifies the installation without side effects?- What the expected output is and whether it should be validated beyond the exit code?
- Whether a failed verification should block all tests or only warn/skip?
|
[FAILED] Pipeline #38459095: 3/17 passed |
kahyunnam
changed the title
There was a problem hiding this comment.
Thanks @kahyunnam! Left a comment about the behavior with jit cache & cubin wheels are not found.
kahyunnam
force-pushed
the
knam/unit-testing-move-compilation-offline
branch
from
0a7f5c7 to
3f6dc1c
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
scripts/task_test_blackwell_kernels.sh (3)
44-45: Validate that relative paths are robust to working directory assumptions.Lines 44–45 construct distribution paths as
../dist/${CUDA_VERSION}/${JIT_ARCH_EFFECTIVE}/.... This assumes the script is invoked from a specific directory (likely the repository root). If the script is called from a different directory, these paths will fail silently or point to unintended locations.Consider either:
- Using
$(dirname "${BASH_SOURCE[0]}")to anchor paths relative to the script location.- Adding explicit validation that
DIST_CUBIN_DIRandDIST_JIT_CACHE_DIRare accessible before attempting installation.- Documenting the expected working directory requirement in a comment.
33-42: Simplify JIT_ARCH mapping logic for clarity.The nested conditional on lines 34–39 is difficult to follow. The logic maps only
12.0to architecture-specific suffixes (12.0afor cu129,12.0fotherwise), while other values pass through unchanged. Consider extracting this into a helper function or adding comments to explain the mapping rules and why12.0is special.
28-29: Add validation or explicit handling for CUDA_VERSION.Line 28 echoes
CUDA_VERSIONbut does not validate that it is set to an expected value (cu128, cu129, cu130, etc.). IfCUDA_VERSIONis unset or malformed, the script will still proceed and construct invalid paths. Consider adding a check to fail fast if the value is unexpected, or document the assumption thatCUDA_VERSIONis always set by the caller.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
scripts/build_flashinfer_jit_cache_whl.sh(1 hunks)scripts/task_test_blackwell_kernels.sh(1 hunks)
⏰ 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). (1)
- GitHub Check: Deploy Docs
🔇 Additional comments (3)
scripts/build_flashinfer_jit_cache_whl.sh (1)
14-15: Clarify rationale for changing MAX_JOBS divisor from architecture-dependent to constant.The change removes the conditional logic and always divides by 8, whereas the original divided by 4 on x86_64. While this simplifies the calculation, it may result in fewer parallel jobs on non-aarch64 systems, potentially increasing build time.
Was this change validated on both architectures? If not, consider testing build times on x86_64 to confirm acceptable performance.
scripts/task_test_blackwell_kernels.sh (2)
65-72: Approve artifact installation and verification flow.The addition of local source installation and verification via
python -m flashinfer show-configis well-structured. Running the verification in/tmpisolates side effects and ensures the installed packages work in a clean environment. The error handling is appropriate for this stage.
51-53: Clarify intent regardingexit 1statements at lines 52 and 60.The current code contains
exit 1at both locations (lines 52 and 60). The review comment references a prior resolution where you stated you "removed the 'exit 1' for both the cubin / jit-cache else logic," but I cannot access the prior conversation to verify this claim.Please confirm:
- Was removing the
exit 1statements intentionally reverted?- Is the current behavior (hard error on missing artifacts) the intended behavior?
- If the intent was to warn and continue with JIT compilation fallback, these statements need to be replaced with warnings.
Without access to the prior review thread, I cannot determine whether this is an oversight or intentional. The developer must clarify the design decision.
4 participants
📌 Description
Download
flashinfer-cubinandflashinfer-jit-cacheto avoid compilation. (Unless the JIT kernel is not in theflashinfer-jit-cache; then it will still JIT compile during test runtime. We could setexport FLASHINFER_DISABLE_JIT = 1to avoid this, but then it will "skip" a lot of tests that use JIT kernels that aren't found inflashinfer-jit-cache.)🔍 Related Issues
Issue was discussed on slack. "Ideally, we would move that compilation off-line which would reduce test time & make kernel hang detection much easier. "
🚀 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Summary by CodeRabbit
Chores
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.