Skip to content

Commit fd03820

Browse files
authored
bugfix: fix cli error when cuda toolkit is not installed (#1905)
<!-- .github/pull_request_template.md --> ## πŸ“Œ Description This is a bugfix to #1904. ## πŸ” Related Issues #1904 ## πŸš€ 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 cc @bbartels
1 parent 940b2f1 commit fd03820

File tree

1 file changed

+14
-2
lines changed

1 file changed

+14
-2
lines changed

β€Žflashinfer/__main__.pyβ€Ž

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""
1616

1717
# flashinfer-cli
18+
import os
1819
import click
1920
from tabulate import tabulate # type: ignore[import-untyped]
2021

@@ -79,13 +80,14 @@ def cli(ctx, download_cubin_flag):
7980
"FLASHINFER_CUDA_ARCH_LIST": current_compilation_context.TARGET_CUDA_ARCHS,
8081
"FLASHINFER_CUDA_VERSION": get_cuda_version(),
8182
"FLASHINFER_CUBINS_REPOSITORY": FLASHINFER_CUBINS_REPOSITORY,
82-
"CUDA_HOME": get_cuda_path(),
8383
"CUDA_VERSION": get_cuda_version(),
8484
}
8585
try:
8686
env_variables["CUDA_HOME"] = get_cuda_path()
87+
found_nvcc = os.path.isfile(os.path.join(env_variables["CUDA_HOME"], "bin", "nvcc"))
8788
except Exception:
88-
env_variables["CUDA_HOME"] = "Not Found"
89+
env_variables["CUDA_HOME"] = ""
90+
found_nvcc = False
8991

9092

9193
@cli.command("show-config")
@@ -124,13 +126,23 @@ def show_config_cmd():
124126
click.secho("=== Torch Version Info ===", fg="yellow")
125127
click.secho("Torch version:", fg="magenta", nl=False)
126128
click.secho(f" {torch.__version__}", fg="cyan")
129+
click.secho("CUDA runtime available:", fg="magenta", nl=False)
130+
if torch.cuda.is_available():
131+
click.secho(" Yes", fg="green")
132+
else:
133+
click.secho(" No", fg="red")
127134
click.secho("", fg="white")
128135

129136
# Section: Environment Variables
130137
click.secho("=== Environment Variables ===", fg="yellow")
131138
for name, value in env_variables.items():
132139
click.secho(f"{name}:", fg="magenta", nl=False)
133140
click.secho(f" {value}", fg="cyan")
141+
click.secho("NVCC found:", fg="magenta", nl=False)
142+
if found_nvcc:
143+
click.secho(" Yes", fg="green")
144+
else:
145+
click.secho(" No", fg="red")
134146
click.secho("", fg="white")
135147

136148
# Section: Artifact path

0 commit comments

Comments
Β (0)